Improve sharded query routing for tuple list

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 @@
 	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 @@
 
 	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 2023 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 (
+	"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)
+	}
+
+	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)
+		}
+		sValue := sqltypes.ProtoToValue(value)
+		var evalErr error
+		idx := 0
+		found := false
+		loopErr := sValue.ForEachValue(func(val sqltypes.Value) {
+			idx++
+			if idx-1 != bv.Index {
+				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 loopErr != nil {
+			return nil, loopErr
+		}
+		if evalErr != nil {
+			return nil, evalErr
+		}
+		if !found {
+			return nil, vterrors.VT13001("value not found in 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/format.go

@@ -152,6 +152,18 @@
 	}
 }
 
+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 @@
 	return false
 }
 
+func (expr *TupleBindVariable) constant() bool {
+	return false
+}
+
 func (expr *Column) constant() bool {
 	return false
 }
@@ -55,6 +59,10 @@
 	return nil
 }
 
+func (expr *TupleBindVariable) simplify(_ *ExpressionEnv) error {
+	return nil
+}
+
 func (expr *Column) simplify(_ *ExpressionEnv) error {
 	return nil
 }

go/vt/vtgate/executor_dml_test.go

@@ -3065,16 +3065,9 @@ func TestDeleteMulti(t *testing.T) {
 
 	testQueryLog(t, executor, logChan, "MarkSavepoint", "SAVEPOINT", "savepoint s1", 8)
 	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
-	testQueryLog(t, executor, logChan, "VindexDelete", "DELETE", "delete from name_user_map where `name` = :name and user_id = :user_id", 1)
 	// select `user`.id, `user`.col from `user` - 8 shard
 	// select 1 from music where music.user_id = 1 and music.col = :user_col - 8 shards
-	// select Id, `name` from `user` where (`user`.id) in ::dm_vals for update - 8 shards
-	// delete from `user` where (`user`.id) in ::dm_vals - 8 shards
-	testQueryLog(t, executor, logChan, "TestExecute", "DELETE", "delete `user` from `user` join music on `user`.col = music.col where music.user_id = 1", 32)
+	// select Id, `name` from `user` where (`user`.id) in ::dm_vals for update - 1 shard
+	// delete from `user` where (`user`.id) in ::dm_vals - 1 shard
+	testQueryLog(t, executor, logChan, "TestExecute", "DELETE", "delete `user` from `user` join music on `user`.col = music.col where music.user_id = 1", 18)
 }

go/vt/vtgate/executor_select_test.go

@@ -1552,6 +1552,75 @@ func TestStreamSelectIN(t *testing.T) {
 	utils.MustMatch(t, wantQueries, sbclookup.Queries)
 }
 
+// TestSelectListArg tests list arg filter with select query
+func TestSelectListArg(t *testing.T) {
+	executor, sbc1, sbc2, _, ctx := createExecutorEnv(t)
+	session := &vtgatepb.Session{
+		TargetString: "@primary",
+	}
+
+	tupleBV := &querypb.BindVariable{
+		Type:   querypb.Type_TUPLE,
+		Values: []*querypb.Value{sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewVarChar("a")))},
+	}
+	bvMap := map[string]*querypb.BindVariable{"vals": tupleBV}
+	_, err := executorExec(ctx, executor, session, "select id from user where (id, col) in ::vals", bvMap)
+	require.NoError(t, err)
+	wantQueries := []*querypb.BoundQuery{{
+		Sql:           "select id from `user` where (id, col) in ::vals",
+		BindVariables: bvMap,
+	}}
+	utils.MustMatch(t, wantQueries, sbc1.Queries)
+	assert.Nil(t, sbc2.Queries, "sbc2.Queries: %+v, want nil", sbc2.Queries)
+
+	sbc1.Queries = nil
+	// get c0-e0 sandbox connection.
+	tbh, err := executor.scatterConn.gateway.hc.GetTabletHealthByAlias(&topodatapb.TabletAlias{
+		Cell: "aa",
+		Uid:  7,
+	})
+	require.NoError(t, err)
+	sbc := tbh.Conn.(*sandboxconn.SandboxConn)
+	sbc.Queries = nil
+
+	_, err = executorExec(ctx, executor, session, "select id from multicol_tbl where (cola, colb) in ::vals", bvMap)
+	require.NoError(t, err)
+
+	wantQueries = []*querypb.BoundQuery{{
+		Sql:           "select id from multicol_tbl where (cola, colb) in ::vals",
+		BindVariables: bvMap,
+	}}
+	utils.MustMatch(t, wantQueries, sbc.Queries)
+	assert.Nil(t, sbc1.Queries, "sbc1.Queries: %+v, want nil", sbc2.Queries)
+	assert.Nil(t, sbc2.Queries, "sbc2.Queries: %+v, want nil", sbc2.Queries)
+
+	tupleBV.Values[0] = sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewInt64(1), sqltypes.NewInt64(42), sqltypes.NewVarChar("a")))
+	sbc.Queries = nil
+	_, err = executorExec(ctx, executor, session, "select id from multicol_tbl where (cola, colx, colb) in ::vals", bvMap)
+	require.NoError(t, err)
+
+	wantQueries = []*querypb.BoundQuery{{
+		Sql:           "select id from multicol_tbl where (cola, colx, colb) in ::vals",
+		BindVariables: bvMap,
+	}}
+	utils.MustMatch(t, wantQueries, sbc.Queries)
+	assert.Nil(t, sbc1.Queries, "sbc1.Queries: %+v, want nil", sbc2.Queries)
+	assert.Nil(t, sbc2.Queries, "sbc2.Queries: %+v, want nil", sbc2.Queries)
+
+	tupleBV.Values[0] = sqltypes.ValueToProto(sqltypes.TestTuple(sqltypes.NewVarChar("a"), sqltypes.NewInt64(42), sqltypes.NewInt64(1)))
+	sbc.Queries = nil
+	_, err = executorExec(ctx, executor, session, "select id from multicol_tbl where (colb, colx, cola) in ::vals", bvMap)
+	require.NoError(t, err)
+
+	wantQueries = []*querypb.BoundQuery{{
+		Sql:           "select id from multicol_tbl where (colb, colx, cola) in ::vals",
+		BindVariables: bvMap,
+	}}
+	utils.MustMatch(t, wantQueries, sbc.Queries)
+	assert.Nil(t, sbc1.Queries, "sbc1.Queries: %+v, want nil", sbc2.Queries)
+	assert.Nil(t, sbc2.Queries, "sbc2.Queries: %+v, want nil", sbc2.Queries)
+}
+
 func createExecutor(ctx context.Context, serv *sandboxTopo, cell string, resolver *Resolver) *Executor {
 	queryLogger := streamlog.New[*logstats.LogStats]("VTGate", queryLogBufferSize)
 	plans := DefaultPlanCache()

go/vt/vtgate/executor_test.go

@@ -889,6 +889,7 @@ func TestExecutorShow(t *testing.T) {
 			buildVarCharRow("TestExecutor", "keyspace_id", "numeric", "", ""),
 			buildVarCharRow("TestExecutor", "krcol_unique_vdx", "keyrange_lookuper_unique", "", ""),
 			buildVarCharRow("TestExecutor", "krcol_vdx", "keyrange_lookuper", "", ""),
+			buildVarCharRow("TestExecutor", "multicol_vdx", "multicol", "column_count=2; column_vindex=xxhash,binary", ""),
 			buildVarCharRow("TestExecutor", "music_user_map", "lookup_hash_unique", "from=music_id; table=music_user_map; to=user_id", "music"),
 			buildVarCharRow("TestExecutor", "name_lastname_keyspace_id_map", "lookup", "from=name,lastname; table=name_lastname_keyspace_id_map; to=keyspace_id", "user2"),
 			buildVarCharRow("TestExecutor", "name_user_map", "lookup_hash", "from=name; table=name_user_map; to=user_id", "user"),

go/vt/vtgate/planbuilder/operators/route.go

@@ -261,7 +261,9 @@
 	opcode func(*vindexes.ColumnVindex) engine.Opcode,
 ) bool {
 	option.ColsSeen[colLoweredName] = true
-	option.ValueExprs = append(option.ValueExprs, valueExpr)
+	if valueExpr != nil {
+		option.ValueExprs = append(option.ValueExprs, valueExpr)
+	}
 	option.Values[indexOfCol] = value
 	option.Predicates[indexOfCol] = node
 	option.Ready = len(option.ColsSeen) == len(colVindex.Columns)