Improve sharded query routing for tuple list (#14892)

Signed-off-by: Harshit Gangal <[email protected]>

go/vt/vtgate/engine/route_test.go

@@ -1716,3 +1716,46 @@ func TestBuildMultiColumnVindexValues(t *testing.T) {
 		})
 	}
 }
+
+// TestSelectTupleMultiCol tests route execution having bind variable with multi column tuple.
+func TestSelectTupleMultiCol(t *testing.T) {
+	vindex, _ := vindexes.CreateVindex("multicol", "", map[string]string{
+		"column_count":  "2",
+		"column_vindex": "hash,binary",
+	})
+
+	sel := NewRoute(
+		MultiEqual,
+		&vindexes.Keyspace{Name: "user", Sharded: true},
+		"select 1 from multicol_tbl where (colb, colx, cola) in ::vals",
+		"select 1 from multicol_tbl where 1 != 1",
+	)
+	sel.Vindex = vindex
+	sel.Values = []evalengine.Expr{
+		&evalengine.TupleBindVariable{Key: "vals", Index: 0},
+		&evalengine.TupleBindVariable{Key: "vals", Index: 1},
+	}
+
+	v1 := sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a"))
+	v2 := sqltypes.TestTuple(sqltypes.NewInt64(4), sqltypes.NewVarChar("b"))
+	tupleBV := &querypb.BindVariable{
+		Type:   querypb.Type_TUPLE,
+		Values: append([]*querypb.Value{sqltypes.ValueToProto(v1)}, sqltypes.ValueToProto(v2)),
+	}
+	vc := &loggingVCursor{
+		shards: []string{"-20", "20-"},
+	}
+	_, err := sel.TryExecute(context.Background(), vc, map[string]*querypb.BindVariable{"vals": tupleBV}, false)
+	require.NoError(t, err)
+	vc.ExpectLog(t, []string{
+		`ResolveDestinationsMultiCol user [[INT64(1) VARCHAR("a")] [INT64(4) VARCHAR("b")]] Destinations:DestinationKeyspaceID(166b40b461),DestinationKeyspaceID(d2fd886762)`,
+		`ExecuteMultiShard user.-20: select 1 from multicol_tbl where (colb, colx, cola) in ::vals {vals: type:TUPLE values:{type:TUPLE value:"\x89\x02\x011\x950\x01a"} values:{type:TUPLE value:"\x89\x02\x014\x950\x01b"}} false false`,
+	})
+
+	vc.Rewind()
+	_, _ = wrapStreamExecute(sel, vc, map[string]*querypb.BindVariable{"vals": tupleBV}, false)
+	vc.ExpectLog(t, []string{
+		`ResolveDestinationsMultiCol user [[INT64(1) VARCHAR("a")] [INT64(4) VARCHAR("b")]] Destinations:DestinationKeyspaceID(166b40b461),DestinationKeyspaceID(d2fd886762)`,
+		`StreamExecuteMulti select 1 from multicol_tbl where (colb, colx, cola) in ::vals user.-20: {vals: type:TUPLE values:{type:TUPLE value:"\x89\x02\x011\x950\x01a"} values:{type:TUPLE value:"\x89\x02\x014\x950\x01b"}} `,
+	})
+}

go/vt/vtgate/evalengine/expr_bvar.go

@@ -65,7 +65,7 @@ func (bv *BindVariable) eval(env *ExpressionEnv) (eval, error) {
 	switch bvar.Type {
 	case sqltypes.Tuple:
 		if bv.Type != sqltypes.Tuple {
-			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "query argument '%s' cannot be a tuple", bv.Key)
+			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "query argument '%s' must be a tuple (is %s)", bv.Key, bvar.Type)
 		}
 
 		tuple := make([]eval, 0, len(bvar.Values))
@@ -80,7 +80,7 @@ func (bv *BindVariable) eval(env *ExpressionEnv) (eval, error) {
 
 	default:
 		if bv.Type == sqltypes.Tuple {
-			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "query argument '%s' must be a tuple (is %s)", bv.Key, bvar.Type)
+			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "query argument '%s' cannot be a tuple", bv.Key)
 		}
 		typ := bvar.Type
 		if bv.typed() {

go/vt/vtgate/evalengine/expr_tuple_bvar.go

@@ -0,0 +1,110 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package evalengine
+
+import (
+	"errors"
+
+	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/sqltypes"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
+)
+
+type (
+	TupleBindVariable struct {
+		Key string
+
+		Index     int
+		Type      sqltypes.Type
+		Collation collations.ID
+
+		// dynamicTypeOffset is set when the type of this bind variable cannot be calculated
+		// at translation time. Since expressions with dynamic types cannot be compiled ahead of time,
+		// compilation will be delayed until the expression is first executed with the bind variables
+		// sent by the user. See: UntypedExpr
+		dynamicTypeOffset int
+	}
+)
+
+var _ IR = (*TupleBindVariable)(nil)
+var _ Expr = (*TupleBindVariable)(nil)
+
+func (bv *TupleBindVariable) IR() IR {
+	return bv
+}
+
+func (bv *TupleBindVariable) IsExpr() {}
+
+// eval implements the expression interface
+func (bv *TupleBindVariable) eval(env *ExpressionEnv) (eval, error) {
+	bvar, err := env.lookupBindVar(bv.Key)
+	if err != nil {
+		return nil, err
+	}
+
+	if bvar.Type != sqltypes.Tuple {
+		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "query argument '%s' must be a tuple (is %s)", bv.Key, bvar.Type.String())
+	}
+
+	tuple := make([]eval, 0, len(bvar.Values))
+	for _, value := range bvar.Values {
+		if value.Type != sqltypes.Tuple {
+			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "result value must be a tuple (is %s)", value.Type.String())
+		}
+		sValue := sqltypes.ProtoToValue(value)
+		var evalErr error
+		idx := 0
+		found := false
+		// looking for a single index on each Tuple Value.
+		loopErr := sValue.ForEachValue(func(val sqltypes.Value) {
+			if found || idx != bv.Index {
+				idx++
+				return
+			}
+			found = true
+			e, err := valueToEval(val, typedCoercionCollation(val.Type(), collations.CollationForType(val.Type(), bv.Collation)))
+			if err != nil {
+				evalErr = err
+				return
+			}
+			tuple = append(tuple, e)
+
+		})
+		if err = errors.Join(loopErr, evalErr); err != nil {
+			return nil, err
+		}
+		if !found {
+			return nil, vterrors.VT13001("value not found in the bind variable")
+		}
+	}
+	return &evalTuple{t: tuple}, nil
+}
+
+// typeof implements the expression interface
+func (bv *TupleBindVariable) typeof(env *ExpressionEnv) (ctype, error) {
+	_, err := env.lookupBindVar(bv.Key)
+	if err != nil {
+		return ctype{}, err
+	}
+
+	return ctype{Type: sqltypes.Tuple}, nil
+}
+
+func (bv *TupleBindVariable) compile(c *compiler) (ctype, error) {
+	return ctype{}, c.unsupported(bv)
+}

go/vt/vtgate/evalengine/expr_tuple_bvar_test.go

@@ -0,0 +1,156 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package evalengine
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+)
+
+// TestTupleBindVarEval tests TupleBindVariable eval function.
+func TestTupleBindVarEval(t *testing.T) {
+	key := "vals"
+	c := &TupleBindVariable{
+		Key:   key,
+		Index: 1,
+	}
+	collation := collations.TypedCollation{
+		Coercibility: collations.CoerceCoercible,
+		Repertoire:   collations.RepertoireUnicode,
+	}
+
+	tcases := []struct {
+		tName string
+		bv    *querypb.BindVariable
+
+		expEval []eval
+		expErr  string
+	}{{
+		tName:  "bind variable not provided",
+		expErr: "query arguments missing for vals",
+	}, {
+		tName:  "bind variable provided - wrong type",
+		bv:     sqltypes.Int64BindVariable(1),
+		expErr: "query argument 'vals' must be a tuple (is INT64)",
+	}, {
+		tName: "bind variable provided",
+		bv: &querypb.BindVariable{
+			Type:   querypb.Type_TUPLE,
+			Values: []*querypb.Value{sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a")))},
+		},
+		expEval: []eval{newEvalText([]byte("a"), collation)},
+	}, {
+		tName: "bind variable provided - multi values",
+		bv: &querypb.BindVariable{
+			Type: querypb.Type_TUPLE,
+			Values: []*querypb.Value{
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a"))),
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(2), sqltypes.NewVarChar("b"))),
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(3), sqltypes.NewVarChar("c"))),
+			},
+		},
+		expEval: []eval{
+			newEvalText([]byte("a"), collation),
+			newEvalText([]byte("b"), collation),
+			newEvalText([]byte("c"), collation)},
+	}}
+
+	for _, tcase := range tcases {
+		t.Run(tcase.tName, func(t *testing.T) {
+			env := &ExpressionEnv{
+				BindVars: make(map[string]*querypb.BindVariable),
+			}
+			if tcase.bv != nil {
+				env.BindVars[key] = tcase.bv
+			}
+
+			res, err := c.eval(env)
+			if tcase.expErr != "" {
+				require.ErrorContains(t, err, tcase.expErr)
+				return
+			}
+			require.Equal(t, sqltypes.Tuple, res.SQLType())
+			resTuple := res.(*evalTuple)
+			require.Len(t, resTuple.t, len(tcase.expEval))
+			for idx, e := range tcase.expEval {
+				require.Equal(t, e, resTuple.t[idx])
+			}
+		})
+	}
+}
+
+// TestTupleBindVarTypeOf tests TupleBindVariable typeOf function.
+func TestTupleBindVarTypeOf(t *testing.T) {
+	key := "vals"
+	c := &TupleBindVariable{
+		Key:   key,
+		Index: 1,
+	}
+
+	tcases := []struct {
+		tName string
+		bv    *querypb.BindVariable
+
+		expErr string
+	}{{
+		tName:  "bind variable not provided",
+		expErr: "query arguments missing for vals",
+	}, {
+		// typeOf does not evaluate the bind variable value
+		tName: "bind variable provided - wrong type",
+		bv:    sqltypes.Int64BindVariable(1),
+	}, {
+		tName: "bind variable provided",
+		bv: &querypb.BindVariable{
+			Type:   querypb.Type_TUPLE,
+			Values: []*querypb.Value{sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a")))},
+		},
+	}, {
+		tName: "bind variable provided - multi values",
+		bv: &querypb.BindVariable{
+			Type: querypb.Type_TUPLE,
+			Values: []*querypb.Value{
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a"))),
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(2), sqltypes.NewVarChar("b"))),
+				sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(3), sqltypes.NewVarChar("c"))),
+			},
+		},
+	}}
+
+	for _, tcase := range tcases {
+		t.Run(tcase.tName, func(t *testing.T) {
+			env := &ExpressionEnv{
+				BindVars: make(map[string]*querypb.BindVariable),
+			}
+			if tcase.bv != nil {
+				env.BindVars[key] = tcase.bv
+			}
+
+			res, err := c.typeof(env)
+			if tcase.expErr != "" {
+				require.ErrorContains(t, err, tcase.expErr)
+				return
+			}
+			require.Equal(t, sqltypes.Tuple, res.Type)
+		})
+	}
+}

go/vt/vtgate/evalengine/format.go

@@ -152,6 +152,18 @@ func (bv *BindVariable) format(buf *sqlparser.TrackedBuffer) {
 	}
 }
 
+func (bv *TupleBindVariable) Format(buf *sqlparser.TrackedBuffer) {
+	bv.format(buf)
+}
+
+func (bv *TupleBindVariable) FormatFast(buf *sqlparser.TrackedBuffer) {
+	bv.format(buf)
+}
+
+func (bv *TupleBindVariable) format(buf *sqlparser.TrackedBuffer) {
+	buf.WriteString(fmt.Sprintf("%s:%d", bv.Key, bv.Index))
+}
+
 func (c *Column) Format(buf *sqlparser.TrackedBuffer) {
 	c.format(buf)
 }

go/vt/vtgate/evalengine/translate_simplify.go

@@ -26,6 +26,10 @@ func (expr *BindVariable) constant() bool {
 	return false
 }
 
+func (expr *TupleBindVariable) constant() bool {
+	return false
+}
+
 func (expr *Column) constant() bool {
 	return false
 }
@@ -55,6 +59,10 @@ func (expr *BindVariable) simplify(_ *ExpressionEnv) error {
 	return nil
 }
 
+func (expr *TupleBindVariable) simplify(_ *ExpressionEnv) error {
+	return nil
+}
+
 func (expr *Column) simplify(_ *ExpressionEnv) error {
 	return nil
 }