Repository: vitessio/vitess
Pull request #9426 (2021-12-20 20:49:40)

Add list support to vindex functions
## Description
This adds support for vindex function batch lookups using the `IN` SQL operator:
```
mysql> select * from customer.hash where id = 100;
+------+-------------+-------------+-----------+------------------+-------+
| id   | keyspace_id | range_start | range_end | hex_keyspace_id  | shard |
+------+-------------+-------------+-----------+------------------+-------+
| 100  | ���V��    | �           |           | 83aab1569cbe1b08 | 80-   |
+------+-------------+-------------+-----------+------------------+-------+
1 row in set (0.00 sec)

mysql> select * from customer.hash where id in (100,1,999);
+------+-------------+-------------+-----------+------------------+-------+
| id   | keyspace_id | range_start | range_end | hex_keyspace_id  | shard |
+------+-------------+-------------+-----------+------------------+-------+
| 100  | ���V��    | �           |           | 83aab1569cbe1b08 | 80-   |
| 1    | k@�J�K�    |             | �         | 166b40b44aba4bd6 | -80   |
| 999  | X�Ο��IH    |             | �         | 58afce9fbfe54948 | -80   |
+------+-------------+-------------+-----------+------------------+-------+
3 rows in set (0.00 sec)

mysql> select * from customer.hash where id like "5";
ERROR 1105 (HY000): unsupported: where clause for vindex function must be of the form id = <val> or id in(<val>,...) (not equality)
```

For more info about the use case, please see: https://github.com/vitessio/vitess/issues/9042

## Related Issue(s)
Fixes: https://github.com/vitessio/vitess/issues/9042


## Checklist
- [x] Should this PR be backported? _**I don't think so**_
- [x] Tests were added
- [x] Documentation is not required -- **_I could not find any documentation to modify for using the vindex functions???_**

## Deployment Notes
<!-- Notes regarding deployment of the contained body of work. These should note any db migrations, etc. -->


Patch:
diff --git a/go/vt/vtgate/engine/vindex_func.go b/go/vt/vtgate/engine/vindex_func.go
index 9dd76f96f6..6873b91750 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 66ce15d1c0..9a64e31192 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 586701903f..e415ee6315 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 2c1188fa6f..ed488590ca 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 3f1867eac7..f27e2f4689 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 31f203d8d7..685e32b747 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"