Repository: vitessio/vitess (Go)

Issue:
Support Batch Lookup With VIndex Functions in Planner
#### Feature Description

Vitess provides a nice method to know where a row would land in the keyspace by allowing you to query the vindex function and pass it an id value to see which shard the row would land on:
```
mysql> select * from customer.hash where id = 5;
+------+-------------+-------------+-----------+------------------+-------+
| id   | keyspace_id | range_start | range_end | hex_keyspace_id  | shard |
+------+-------------+-------------+-----------+------------------+-------+
| 5    | p�<�
             �z    |             | �         | 70bb023c810ca87a | -80   |
+------+-------------+-------------+-----------+------------------+-------+
1 row in set (0.00 sec)
```

However, you must perform this check one id at a time. This makes it less convenient when you are planing to do large bulk inserts / data loads:
```
mysql> select id, hex_keyspace_id, shard from customer.hash where id IN (5,9999);
ERROR 1105 (HY000): unsupported: where clause for vindex function must be of the form id = <val> (not equality)

mysql> select id, hex_keyspace_id, shard from customer.hash where id = 5;
+------+------------------+-------+
| id   | hex_keyspace_id  | shard |
+------+------------------+-------+
| 5    | 70bb023c810ca87a | -80   |
+------+------------------+-------+
1 row in set (0.00 sec)

mysql> select id, hex_keyspace_id, shard from customer.hash where id = 9999;
+------+------------------+-------+
| id   | hex_keyspace_id  | shard |
+------+------------------+-------+
| 9999 | 9a49bef394017d2d | 80-   |
+------+------------------+-------+
1 row in set (0.00 sec)
```

#### Use Case(s)

It's common to have bulk data copy/load pipelines (ETL or other asynchronous data pipelines). In this pipeline you will often perform transformations to support loading this data in the most efficient way knowing the details of the target database. With Vitess, it would be nice to check the landing `keyspace_id`/`shard` for rows in bulk in order to more efficiently organize data in the pipeline so that it's grouped by shard in order to load the data as quickly as possible in Vitess with limited impact/load on the database (e.g. generating multi-row INSERTs for each shard).



Discussion:

Patch:
diff --git a/go/vt/vtgate/engine/vindex_func.go b/go/vt/vtgate/engine/vindex_func.go
index 9dd76f96f6b..6873b91750b 100644
--- a/go/vt/vtgate/engine/vindex_func.go
+++ b/go/vt/vtgate/engine/vindex_func.go
@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -28,6 +29,8 @@ import (
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/proto/vtrpc"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 )
 
 var _ Primitive = (*VindexFunc)(nil)
@@ -113,57 +116,75 @@ func (vf *VindexFunc) mapVindex(vcursor VCursor, bindVars map[string]*querypb.Bi
 	if err != nil {
 		return nil, err
 	}
-	value := k.Value()
-	vkey, err := evalengine.Cast(value, sqltypes.VarBinary)
-	if err != nil {
-		return nil, err
+	var values []sqltypes.Value
+	if k.Value().Type() == querypb.Type_TUPLE {
+		values = k.TupleValues()
+	} else {
+		values = append(values, k.Value())
 	}
 	result := &sqltypes.Result{
 		Fields: vf.Fields,
 	}
-
-	destinations, err := vf.Vindex.Map(vcursor, []sqltypes.Value{value})
-	if err != nil {
-		return nil, err
-	}
-	switch d := destinations[0].(type) {
-	case key.DestinationKeyRange:
-		if d.KeyRange != nil {
-			result.Rows = append(result.Rows, vf.buildRow(vkey, nil, d.KeyRange))
+	for _, value := range values {
+		vkey, err := evalengine.Cast(value, sqltypes.VarBinary)
+		if err != nil {
+			return nil, err
 		}
-	case key.DestinationKeyspaceID:
-		if len(d) > 0 {
-			if vcursor != nil {
-				resolvedShards, _, err := vcursor.ResolveDestinations(vcursor.GetKeyspace(), nil, []key.Destination{d})
-				if err != nil {
-					return nil, err
-				}
-				kr, err := key.ParseShardingSpec(resolvedShards[0].Target.Shard)
+		destinations, err := vf.Vindex.Map(vcursor, []sqltypes.Value{value})
+		if err != nil {
+			return nil, err
+		}
+		switch d := destinations[0].(type) {
+		case key.DestinationKeyRange:
+			if d.KeyRange != nil {
+				row, err := vf.buildRow(vkey, nil, d.KeyRange)
 				if err != nil {
-					return nil, err
+					return result, err
 				}
-				result.Rows = [][]sqltypes.Value{
-					vf.buildRow(vkey, d, kr[0]),
+				result.Rows = append(result.Rows, row)
+			}
+		case key.DestinationKeyspaceID:
+			if len(d) > 0 {
+				if vcursor != nil {
+					resolvedShards, _, err := vcursor.ResolveDestinations(vcursor.GetKeyspace(), nil, []key.Destination{d})
+					if err != nil {
+						return nil, err
+					}
+					kr, err := key.ParseShardingSpec(resolvedShards[0].Target.Shard)
+					if err != nil {
+						return nil, err
+					}
+					row, err := vf.buildRow(vkey, d, kr[0])
+					if err != nil {
+						return result, err
+					}
+					result.Rows = append(result.Rows, row)
+				} else {
+					row, err := vf.buildRow(vkey, d, nil)
+					if err != nil {
+						return result, err
+					}
+					result.Rows = append(result.Rows, row)
 				}
-			} else {
-				result.Rows = [][]sqltypes.Value{
-					vf.buildRow(vkey, d, nil),
+			}
+		case key.DestinationKeyspaceIDs:
+			for _, ksid := range d {
+				row, err := vf.buildRow(vkey, ksid, nil)
+				if err != nil {
+					return result, err
 				}
+				result.Rows = append(result.Rows, row)
 			}
+		case key.DestinationNone:
+			// Nothing to do.
+		default:
+			return result, vterrors.NewErrorf(vtrpcpb.Code_INTERNAL, vterrors.WrongTypeForVar, "unexpected destination type: %T", d)
 		}
-	case key.DestinationKeyspaceIDs:
-		for _, ksid := range d {
-			result.Rows = append(result.Rows, vf.buildRow(vkey, ksid, nil))
-		}
-	case key.DestinationNone:
-		// Nothing to do.
-	default:
-		panic("unexpected")
 	}
 	return result, nil
 }
 
-func (vf *VindexFunc) buildRow(id sqltypes.Value, ksid []byte, kr *topodatapb.KeyRange) []sqltypes.Value {
+func (vf *VindexFunc) buildRow(id sqltypes.Value, ksid []byte, kr *topodatapb.KeyRange) ([]sqltypes.Value, error) {
 	row := make([]sqltypes.Value, 0, len(vf.Fields))
 	for _, col := range vf.Cols {
 		switch col {
@@ -200,10 +221,10 @@ func (vf *VindexFunc) buildRow(id sqltypes.Value, ksid []byte, kr *topodatapb.Ke
 				row = append(row, sqltypes.NULL)
 			}
 		default:
-			panic("BUG: unexpected column number")
+			return row, vterrors.NewErrorf(vtrpc.Code_OUT_OF_RANGE, vterrors.BadFieldError, "column %v out of range", col)
 		}
 	}
-	return row
+	return row, nil
 }
 
 func (vf *VindexFunc) description() PrimitiveDescription {
diff --git a/go/vt/vtgate/planbuilder/abstract/vindex.go b/go/vt/vtgate/planbuilder/abstract/vindex.go
index 66ce15d1c05..9a64e31192c 100644
--- a/go/vt/vtgate/planbuilder/abstract/vindex.go
+++ b/go/vt/vtgate/planbuilder/abstract/vindex.go
@@ -51,7 +51,7 @@ func (v *Vindex) TableID() semantics.TableSet {
 	return v.Table.TableID
 }
 
-const vindexUnsupported = "unsupported: where clause for vindex function must be of the form id = <val>"
+const vindexUnsupported = "unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...)"
 
 // PushPredicate implements the Operator interface
 func (v *Vindex) PushPredicate(expr sqlparser.Expr, semTable *semantics.SemTable) error {
@@ -70,7 +70,7 @@ func (v *Vindex) PushPredicate(expr sqlparser.Expr, semTable *semantics.SemTable
 		if !ok {
 			return vterrors.Errorf(vtrpcpb.Code_INTERNAL, vindexUnsupported+" (not a comparison)")
 		}
-		if comparison.Operator != sqlparser.EqualOp {
+		if comparison.Operator != sqlparser.EqualOp && comparison.Operator != sqlparser.InOp {
 			return vterrors.Errorf(vtrpcpb.Code_INTERNAL, vindexUnsupported+" (not equality)")
 		}
 		colname, ok := comparison.Left.(*sqlparser.ColName)
@@ -82,11 +82,12 @@ func (v *Vindex) PushPredicate(expr sqlparser.Expr, semTable *semantics.SemTable
 		}
 
 		// check RHS
-		if !sqlparser.IsValue(comparison.Right) {
+		var err error
+		if sqlparser.IsValue(comparison.Right) || sqlparser.IsSimpleTuple(comparison.Right) {
+			v.Value = comparison.Right
+		} else {
 			return vterrors.Errorf(vtrpcpb.Code_INTERNAL, vindexUnsupported+" (rhs is not a value)")
 		}
-		var err error
-		v.Value = comparison.Right
 		if err != nil {
 			return vterrors.Errorf(vtrpcpb.Code_INTERNAL, vindexUnsupported+": %v", err)
 		}
@@ -104,7 +105,7 @@ func (v *Vindex) UnsolvedPredicates(*semantics.SemTable) []sqlparser.Expr {
 // CheckValid implements the Operator interface
 func (v *Vindex) CheckValid() error {
 	if len(v.Table.Predicates) == 0 {
-		return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: where clause for vindex function must be of the form id = <val> (where clause missing)")
+		return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (where clause missing)")
 	}
 
 	return nil
diff --git a/go/vt/vtgate/planbuilder/filtering.go b/go/vt/vtgate/planbuilder/filtering.go
index 586701903f1..e415ee63152 100644
--- a/go/vt/vtgate/planbuilder/filtering.go
+++ b/go/vt/vtgate/planbuilder/filtering.go
@@ -86,35 +86,36 @@ func planFilter(pb *primitiveBuilder, input logicalPlan, filter sqlparser.Expr,
 
 func filterVindexFunc(node *vindexFunc, filter sqlparser.Expr) (logicalPlan, error) {
 	if node.eVindexFunc.Opcode != engine.VindexNone {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (multiple filters)")
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (multiple filters)")
 	}
 
 	// Check LHS.
 	comparison, ok := filter.(*sqlparser.ComparisonExpr)
 	if !ok {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (not a comparison)")
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (not a comparison)")
 	}
-	if comparison.Operator != sqlparser.EqualOp {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (not equality)")
+	if comparison.Operator != sqlparser.EqualOp && comparison.Operator != sqlparser.InOp {
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (not equality)")
 	}
 	colname, ok := comparison.Left.(*sqlparser.ColName)
 	if !ok {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (lhs is not a column)")
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (lhs is not a column)")
 	}
 	if !colname.Name.EqualString("id") {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (lhs is not id)")
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (lhs is not id)")
 	}
 
 	// Check RHS.
 	// We have to check before calling NewPlanValue because NewPlanValue allows lists also.
-	if !sqlparser.IsValue(comparison.Right) {
-		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (rhs is not a value)")
+	if !sqlparser.IsValue(comparison.Right) && !sqlparser.IsSimpleTuple(comparison.Right) {
+		return nil, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (rhs is not a value)")
 	}
 	var err error
 	node.eVindexFunc.Value, err = evalengine.Convert(comparison.Right, semantics.EmptySemTable())
 	if err != nil {
-		return nil, fmt.Errorf("unsupported: where clause for vindex function must be of the form id = <val>: %v", err)
+		return nil, fmt.Errorf("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...): %v", err)
 	}
+
 	node.eVindexFunc.Opcode = engine.VindexMap
 	return node, nil
 }
diff --git a/go/vt/vtgate/planbuilder/project.go b/go/vt/vtgate/planbuilder/project.go
index 2c1188fa6f1..ed488590ca9 100644
--- a/go/vt/vtgate/planbuilder/project.go
+++ b/go/vt/vtgate/planbuilder/project.go
@@ -151,7 +151,7 @@ func planProjection(pb *primitiveBuilder, in logicalPlan, expr *sqlparser.Aliase
 		// Catch the case where no where clause was specified. If so, the opcode
 		// won't be set.
 		if node.eVindexFunc.Opcode == engine.VindexNone {
-			return nil, nil, 0, errors.New("unsupported: where clause for vindex function must be of the form id = <val> (where clause missing)")
+			return nil, nil, 0, errors.New("unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (where clause missing)")
 		}
 		col, ok := expr.Expr.(*sqlparser.ColName)
 		if !ok {


Test patch:
diff --git a/go/vt/vtgate/engine/vindex_func_test.go b/go/vt/vtgate/engine/vindex_func_test.go
index 3f1867eac7b..f27e2f4689e 100644
--- a/go/vt/vtgate/engine/vindex_func_test.go
+++ b/go/vt/vtgate/engine/vindex_func_test.go
@@ -20,6 +20,8 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
 
 	"vitess.io/vitess/go/sqltypes"
@@ -96,22 +98,16 @@ func TestVindexFuncMap(t *testing.T) {
 	// Unique Vindex returning 0 rows.
 	vf := testVindexFunc(&uvindex{})
 	got, err := vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want := &sqltypes.Result{
 		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
-	}
+	require.Equal(t, got, want)
 
 	// Unique Vindex returning 1 row.
 	vf = testVindexFunc(&uvindex{matchid: true})
 	got, err = vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want = sqltypes.MakeTestResult(
 		sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		"1|foo|||666f6f",
@@ -120,16 +116,34 @@ func TestVindexFuncMap(t *testing.T) {
 		row[2] = sqltypes.NULL
 		row[3] = sqltypes.NULL
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
+	require.Equal(t, got, want)
+
+	// Unique Vindex returning 3 rows
+	vf = &VindexFunc{
+		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
+		Cols:   []int{0, 1, 2, 3, 4},
+		Opcode: VindexMap,
+		Vindex: &uvindex{matchid: true},
+		Value:  evalengine.TupleExpr{evalengine.NewLiteralInt(1), evalengine.NewLiteralInt(2), evalengine.NewLiteralInt(3)},
+	}
+	got, err = vf.TryExecute(nil, nil, false)
+	require.NoError(t, err)
+	want = sqltypes.MakeTestResult(
+		sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
+		"1|foo|||666f6f",
+		"2|foo|||666f6f",
+		"3|foo|||666f6f",
+	)
+	for _, row := range want.Rows {
+		row[2] = sqltypes.NULL
+		row[3] = sqltypes.NULL
 	}
+	require.Equal(t, got, want)
 
 	// Unique Vindex returning keyrange.
 	vf = testVindexFunc(&uvindex{matchkr: true})
 	got, err = vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want = &sqltypes.Result{
 		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		Rows: [][]sqltypes.Value{{
@@ -141,29 +155,21 @@ func TestVindexFuncMap(t *testing.T) {
 		}},
 		RowsAffected: 0,
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
-	}
+	require.Equal(t, got, want)
 
 	// NonUnique Vindex returning 0 rows.
 	vf = testVindexFunc(&nvindex{})
 	got, err = vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want = &sqltypes.Result{
 		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
-	}
+	require.Equal(t, got, want)
 
 	// NonUnique Vindex returning 2 rows.
 	vf = testVindexFunc(&nvindex{matchid: true})
 	got, err = vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want = sqltypes.MakeTestResult(
 		sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		"1|foo|||666f6f",
@@ -174,16 +180,12 @@ func TestVindexFuncMap(t *testing.T) {
 		row[2] = sqltypes.NULL
 		row[3] = sqltypes.NULL
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
-	}
+	require.Equal(t, got, want)
 
 	// NonUnique Vindex returning keyrange
 	vf = testVindexFunc(&nvindex{matchkr: true})
 	got, err = vf.TryExecute(nil, nil, false)
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	want = &sqltypes.Result{
 		Fields: sqltypes.MakeTestFields("id|keyspace_id|hex(keyspace_id)|range_start|range_end", "varbinary|varbinary|varbinary|varbinary|varbinary"),
 		Rows: [][]sqltypes.Value{{
@@ -195,9 +197,7 @@ func TestVindexFuncMap(t *testing.T) {
 		}},
 		RowsAffected: 0,
 	}
-	if !reflect.DeepEqual(got, want) {
-		t.Errorf("Execute(Map, uvindex(none)):\n%v, want\n%v", got, want)
-	}
+	require.Equal(t, got, want)
 }
 
 func TestVindexFuncStreamExecute(t *testing.T) {
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
index 31f203d8d77..685e32b747a 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
@@ -329,37 +329,37 @@ Gen4 plan same as above
 Gen4 plan same as above
 
 "select keyspace_id from user_index where id = 1 and id = 2"
-"unsupported: where clause for vindex function must be of the form id = <val> (multiple filters)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (multiple filters)"
 Gen4 plan same as above
 
 "select keyspace_id from user_index where func(id)"
-"unsupported: where clause for vindex function must be of the form id = <val> (not a comparison)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (not a comparison)"
 Gen4 plan same as above
 
 "select keyspace_id from user_index where id > 1"
-"unsupported: where clause for vindex function must be of the form id = <val> (not equality)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (not equality)"
 Gen4 plan same as above
 
 "select keyspace_id from user_index where 1 = id"
-"unsupported: where clause for vindex function must be of the form id = <val> (lhs is not a column)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (lhs is not a column)"
 Gen4 plan same as above
 
 "select keyspace_id from user_index where keyspace_id = 1"
-"unsupported: where clause for vindex function must be of the form id = <val> (lhs is not id)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (lhs is not id)"
 Gen4 plan same as above
 
 "select keyspace_id from user_index where id = id+1"
-"unsupported: where clause for vindex function must be of the form id = <val> (rhs is not a value)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (rhs is not a value)"
 Gen4 plan same as above
 
 # vindex func without where condition
 "select keyspace_id from user_index"
-"unsupported: where clause for vindex function must be of the form id = <val> (where clause missing)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (where clause missing)"
 Gen4 plan same as above
 
 # vindex func in subquery without where
 "select id from user where exists(select keyspace_id from user_index)"
-"unsupported: where clause for vindex function must be of the form id = <val> (where clause missing)"
+"unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (where clause missing)"
 Gen4 plan same as above
 
 "select func(keyspace_id) from user_index where id = :id"