Repository: vitessio/vitess
Pull request #10151 (2022-04-27 08:46:54)

feat: make gen4 plan subsharding queries better
## Description
The gen4 planner was treating partial vindex lookups as equivalent to vindex lookups with all columns accounted for.
This lead to bad plans when some of the column comparisons were `IN` comparisons.

## Related Issue(s)

## Checklist
-   [ ] "Backport me!" label has been added if this change should be backported
-   [x] Tests were added or are not required
-   [x] Documentation was added or is not required



Patch:
diff --git a/go/vt/vtgate/engine/delete.go b/go/vt/vtgate/engine/delete.go
index 8869e3b313..12d8cbb8af 100644
--- a/go/vt/vtgate/engine/delete.go
+++ b/go/vt/vtgate/engine/delete.go
@@ -78,7 +78,7 @@ func (del *Delete) TryExecute(vcursor VCursor, bindVars map[string]*querypb.Bind
 	switch del.Opcode {
 	case Unsharded:
 		return del.execUnsharded(vcursor, bindVars, rss)
-	case Equal, IN, Scatter, ByDestination:
+	case Equal, IN, Scatter, ByDestination, SubShard:
 		return del.execMultiDestination(vcursor, bindVars, rss, del.deleteVindexEntries)
 	default:
 		// Unreachable.
diff --git a/go/vt/vtgate/engine/routing.go b/go/vt/vtgate/engine/routing.go
index 28d038f94c..b382e46d30 100644
--- a/go/vt/vtgate/engine/routing.go
+++ b/go/vt/vtgate/engine/routing.go
@@ -68,8 +68,8 @@ const (
 	// Is used when the query explicitly sets a target destination:
 	// in the clause e.g: UPDATE `keyspace[-]`.x1 SET foo=1
 	ByDestination
-	// NumOpcodes is the number of opcodes
-	NumOpcodes
+	// SubShard is for when we are missing one or more columns from a composite vindex
+	SubShard
 )
 
 var opName = map[Opcode]string{
@@ -84,6 +84,7 @@ var opName = map[Opcode]string{
 	Reference:     "Reference",
 	None:          "None",
 	ByDestination: "ByDestination",
+	SubShard:      "SubShard",
 }
 
 // MarshalJSON serializes the Opcode as a JSON string.
@@ -133,7 +134,7 @@ func (rp *RoutingParameters) findRoute(vcursor VCursor, bindVars map[string]*que
 		return rp.byDestination(vcursor, bindVars, key.DestinationAllShards{})
 	case ByDestination:
 		return rp.byDestination(vcursor, bindVars, rp.TargetDestination)
-	case Equal, EqualUnique:
+	case Equal, EqualUnique, SubShard:
 		switch rp.Vindex.(type) {
 		case vindexes.MultiColumn:
 			return rp.equalMultiCol(vcursor, bindVars)
diff --git a/go/vt/vtgate/engine/update.go b/go/vt/vtgate/engine/update.go
index 7f6e6d31e4..815034a659 100644
--- a/go/vt/vtgate/engine/update.go
+++ b/go/vt/vtgate/engine/update.go
@@ -88,7 +88,7 @@ func (upd *Update) TryExecute(vcursor VCursor, bindVars map[string]*querypb.Bind
 	switch upd.Opcode {
 	case Unsharded:
 		return upd.execUnsharded(vcursor, bindVars, rss)
-	case Equal, EqualUnique, IN, Scatter, ByDestination:
+	case Equal, EqualUnique, IN, Scatter, ByDestination, SubShard:
 		return upd.execMultiDestination(vcursor, bindVars, rss, upd.updateVindexEntries)
 	default:
 		// Unreachable.
diff --git a/go/vt/vtgate/planbuilder/insert.go b/go/vt/vtgate/planbuilder/insert.go
index d6a6573f7c..fa39706647 100644
--- a/go/vt/vtgate/planbuilder/insert.go
+++ b/go/vt/vtgate/planbuilder/insert.go
@@ -295,7 +295,7 @@ func applyCommentDirectives(ins *sqlparser.Insert, eins *engine.Insert) {
 
 func getColVindexes(allColVindexes []*vindexes.ColumnVindex) (colVindexes []*vindexes.ColumnVindex) {
 	for _, colVindex := range allColVindexes {
-		if colVindex.IgnoreInDML() {
+		if colVindex.IsPartialVindex() {
 			continue
 		}
 		colVindexes = append(colVindexes, colVindex)
diff --git a/go/vt/vtgate/planbuilder/physical/route.go b/go/vt/vtgate/planbuilder/physical/route.go
index 3cccfd60d6..bc529c1147 100644
--- a/go/vt/vtgate/planbuilder/physical/route.go
+++ b/go/vt/vtgate/planbuilder/physical/route.go
@@ -355,7 +355,7 @@ func (r *Route) haveMatchingVindex(
 			}
 			v.Options = append(v.Options, newOption...)
 
-			// multi column vindex - just always add as new option for furince we do not have one already
+			// multi column vindex - just always add as new option
 			option := createOption(v.ColVindex, vfunc)
 			optionReady := option.updateWithNewColumn(colLoweredName, valueExpr, indexOfCol, value, node, v.ColVindex, opcode)
 			if optionReady {
@@ -635,6 +635,9 @@ func tupleAccess(expr sqlparser.Expr, coordinates []int) sqlparser.Expr {
 }
 
 func equalOrEqualUnique(vindex *vindexes.ColumnVindex) engine.Opcode {
+	if vindex.IsPartialVindex() {
+		return engine.SubShard
+	}
 	if vindex.IsUnique() {
 		return engine.EqualUnique
 	}
diff --git a/go/vt/vtgate/vindexes/vschema.go b/go/vt/vtgate/vindexes/vschema.go
index 186f9920a4..e424006422 100644
--- a/go/vt/vtgate/vindexes/vschema.go
+++ b/go/vt/vtgate/vindexes/vschema.go
@@ -106,14 +106,14 @@ type Keyspace struct {
 
 // ColumnVindex contains the index info for each index of a table.
 type ColumnVindex struct {
-	Columns     []sqlparser.ColIdent `json:"columns"`
-	Type        string               `json:"type"`
-	Name        string               `json:"name"`
-	Owned       bool                 `json:"owned,omitempty"`
-	Vindex      Vindex               `json:"vindex"`
-	isUnique    bool
-	cost        int
-	ignoreInDML bool
+	Columns  []sqlparser.ColIdent `json:"columns"`
+	Type     string               `json:"type"`
+	Name     string               `json:"name"`
+	Owned    bool                 `json:"owned,omitempty"`
+	Vindex   Vindex               `json:"vindex"`
+	isUnique bool
+	cost     int
+	partial  bool
 }
 
 // IsUnique is used to tell whether the ColumnVindex
@@ -129,9 +129,9 @@ func (c *ColumnVindex) Cost() int {
 	return c.cost
 }
 
-// IgnoreInDML is used to let planner and engine know that they need to be ignored for dml queries.
-func (c *ColumnVindex) IgnoreInDML() bool {
-	return c.ignoreInDML
+// IsPartialVindex is used to let planner and engine know that this is a composite vindex missing one or more columns
+func (c *ColumnVindex) IsPartialVindex() bool {
+	return c.partial
 }
 
 // Column describes a column.
@@ -374,13 +374,13 @@ func buildTables(ks *vschemapb.Keyspace, vschema *VSchema, ksvschema *KeyspaceSc
 				columnSubset := columns[:i]
 				cost++
 				columnVindex = &ColumnVindex{
-					Columns:     columnSubset,
-					Type:        vindexInfo.Type,
-					Name:        ind.Name,
-					Owned:       owned,
-					Vindex:      vindex,
-					cost:        cost,
-					ignoreInDML: true,
+					Columns: columnSubset,
+					Type:    vindexInfo.Type,
+					Name:    ind.Name,
+					Owned:   owned,
+					Vindex:  vindex,
+					cost:    cost,
+					partial: true,
 				}
 				t.ColumnVindexes = append(t.ColumnVindexes, columnVindex)
 			}


Test patch:
diff --git a/go/vt/vtgate/engine/delete_test.go b/go/vt/vtgate/engine/delete_test.go
index 59e980f605..99a07c804c 100644
--- a/go/vt/vtgate/engine/delete_test.go
+++ b/go/vt/vtgate/engine/delete_test.go
@@ -547,7 +547,7 @@ func TestDeleteEqualSubshard(t *testing.T) {
 	del := &Delete{
 		DML: &DML{
 			RoutingParameters: &RoutingParameters{
-				Opcode: Equal,
+				Opcode: SubShard,
 				Keyspace: &vindexes.Keyspace{
 					Name:    "ks",
 					Sharded: true,
diff --git a/go/vt/vtgate/engine/insert_test.go b/go/vt/vtgate/engine/insert_test.go
index 5ec28be1d4..b5a723df46 100644
--- a/go/vt/vtgate/engine/insert_test.go
+++ b/go/vt/vtgate/engine/insert_test.go
@@ -712,7 +712,7 @@ func TestInsertShardedGeo(t *testing.T) {
 		" suffix",
 	)
 	for _, colVindex := range ks.Tables["t1"].ColumnVindexes {
-		if colVindex.IgnoreInDML() {
+		if colVindex.IsPartialVindex() {
 			continue
 		}
 		ins.ColVindexes = append(ins.ColVindexes, colVindex)
diff --git a/go/vt/vtgate/engine/update_test.go b/go/vt/vtgate/engine/update_test.go
index 7dbd2fe689..0cbd3b4949 100644
--- a/go/vt/vtgate/engine/update_test.go
+++ b/go/vt/vtgate/engine/update_test.go
@@ -884,7 +884,7 @@ func TestUpdateEqualSubshard(t *testing.T) {
 	upd := &Update{
 		DML: &DML{
 			RoutingParameters: &RoutingParameters{
-				Opcode: Equal,
+				Opcode: SubShard,
 				Keyspace: &vindexes.Keyspace{
 					Name:    "ks",
 					Sharded: true,
diff --git a/go/vt/vtgate/planbuilder/route_test.go b/go/vt/vtgate/planbuilder/route_test.go
index cd4f229e48..eb5a2310d8 100644
--- a/go/vt/vtgate/planbuilder/route_test.go
+++ b/go/vt/vtgate/planbuilder/route_test.go
@@ -28,22 +28,26 @@ import (
 )
 
 /*
+
+This test file only tests the V3 planner. It does not test the Subshard opcode
+
 For easy reference, opcodes are:
-	SelectUnsharded   0
-	SelectEqualUnique 1
-	SelectEqual       2
-	SelectIN          3
-	SelectMultiEqual  4
-	SelectScatter     5
-	SelectNext        6
-	SelectDBA         7
-	SelectReference   8
-	SelectNone        9
-	NumRouteOpcodes   10
+	Unsharded   	 0
+	EqualUnique 	 1
+	Equal       	 2
+	IN          	 3
+	MultiEqual  	 4
+	Scatter     	 5
+	Next        	 6
+	DBA         	 7
+	Reference   	 8
+	None        	 9
+	Subshard         10 <- not covered
+	NumRouteOpcodes  11
 */
 
 func TestJoinCanMerge(t *testing.T) {
-	testcases := [engine.NumOpcodes][engine.NumOpcodes]bool{
+	testcases := [][]bool{
 		{true, false, false, false, false, false, false, false, true, false, false},
 		{false, true, false, false, false, false, false, false, true, false, false},
 		{false, false, false, false, false, false, false, false, true, false, false},
@@ -60,7 +64,7 @@ func TestJoinCanMerge(t *testing.T) {
 	ks := &vindexes.Keyspace{}
 	for left, vals := range testcases {
 		for right, val := range vals {
-			name := fmt.Sprintf("%d:%d", left, right)
+			name := fmt.Sprintf("%s:%s", engine.Opcode(left).String(), engine.Opcode(right).String())
 			t.Run(name, func(t *testing.T) {
 				lRoute := &route{
 					// Setting condition will make SelectEqualUnique match itself.
@@ -81,7 +85,7 @@ func TestJoinCanMerge(t *testing.T) {
 }
 
 func TestSubqueryCanMerge(t *testing.T) {
-	testcases := [engine.NumOpcodes][engine.NumOpcodes]bool{
+	testcases := [][]bool{
 		{true, false, false, false, false, false, false, false, true, false, false},
 		{false, false, false, false, false, false, false, false, true, false, false},
 		{false, false, false, false, false, false, false, false, true, false, false},
@@ -111,7 +115,7 @@ func TestSubqueryCanMerge(t *testing.T) {
 }
 
 func TestUnionCanMerge(t *testing.T) {
-	testcases := [engine.NumOpcodes][engine.NumOpcodes]bool{
+	testcases := [][]bool{
 		{true, false, false, false, false, false, false, false, false, false, false},
 		{false, false, false, false, false, false, false, false, false, false, false},
 		{false, false, false, false, false, false, false, false, false, false, false},
@@ -124,6 +128,7 @@ func TestUnionCanMerge(t *testing.T) {
 		{false, false, false, false, false, false, false, false, false, false, false},
 		{false, false, false, false, false, false, false, false, false, false, false},
 	}
+
 	ks := &vindexes.Keyspace{}
 	lRoute := &route{}
 	rRoute := &route{}
diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.txt b/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
index 665d9f24c6..4480b0af93 100644
--- a/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
@@ -3309,26 +3309,7 @@ Gen4 plan same as above
     "Vindex": "multicolIdx"
   }
 }
-{
-  "QueryType": "UPDATE",
-  "Original": "update multicol_tbl set x = 1 where colb IN (1,2) and cola = 1",
-  "Instructions": {
-    "OperatorType": "Update",
-    "Variant": "Equal",
-    "Keyspace": {
-      "Name": "user",
-      "Sharded": true
-    },
-    "TargetTabletType": "PRIMARY",
-    "MultiShardAutocommit": false,
-    "Query": "update multicol_tbl set x = 1 where colb in (1, 2) and cola = 1",
-    "Table": "multicol_tbl",
-    "Values": [
-      "INT64(1)"
-    ],
-    "Vindex": "multicolIdx"
-  }
-}
+Gen4 plan same as above
 
 # update with a multicol vindex using an IN clause
 "update multicol_tbl set x = 1 where colb IN (1,2) and cola IN (3,4)"
@@ -3570,7 +3551,26 @@ Gen4 plan same as above
     "Vindex": "multicolIdx"
   }
 }
-Gen4 plan same as above
+{
+  "QueryType": "UPDATE",
+  "Original": "update multicol_tbl set x = 42 where cola = 1",
+  "Instructions": {
+    "OperatorType": "Update",
+    "Variant": "SubShard",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "TargetTabletType": "PRIMARY",
+    "MultiShardAutocommit": false,
+    "Query": "update multicol_tbl set x = 42 where cola = 1",
+    "Table": "multicol_tbl",
+    "Values": [
+      "INT64(1)"
+    ],
+    "Vindex": "multicolIdx"
+  }
+}
 
 # update with routing using subsharding column on lookup vindex
 "update multicol_tbl set name = 'bar' where cola = 1"
@@ -3600,7 +3600,32 @@ Gen4 plan same as above
     "Vindex": "multicolIdx"
   }
 }
-Gen4 plan same as above
+{
+  "QueryType": "UPDATE",
+  "Original": "update multicol_tbl set name = 'bar' where cola = 1",
+  "Instructions": {
+    "OperatorType": "Update",
+    "Variant": "SubShard",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "TargetTabletType": "PRIMARY",
+    "ChangedVindexValues": [
+      "name_muticoltbl_map:4"
+    ],
+    "KsidLength": 2,
+    "KsidVindex": "multicolIdx",
+    "MultiShardAutocommit": false,
+    "OwnedVindexQuery": "select cola, colb, colc, `name`, `name` = 'bar' from multicol_tbl where cola = 1 for update",
+    "Query": "update multicol_tbl set `name` = 'bar' where cola = 1",
+    "Table": "multicol_tbl",
+    "Values": [
+      "INT64(1)"
+    ],
+    "Vindex": "multicolIdx"
+  }
+}
 
 # update with routing using subsharding column with in query
 "update multicol_tbl set name = 'bar' where cola in (1,2)"
@@ -3654,7 +3679,26 @@ Gen4 plan same as above
     "Vindex": "multicolIdx"
   }
 }
-Gen4 plan same as above
+{
+  "QueryType": "UPDATE",
+  "Original": "update multicol_tbl set x = 1 where name = 'foo' and cola = 2",
+  "Instructions": {
+    "OperatorType": "Update",
+    "Variant": "Equal",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "TargetTabletType": "PRIMARY",
+    "MultiShardAutocommit": false,
+    "Query": "update multicol_tbl set x = 1 where `name` = 'foo' and cola = 2",
+    "Table": "multicol_tbl",
+    "Values": [
+      "VARCHAR(\"foo\")"
+    ],
+    "Vindex": "name_muticoltbl_map"
+  }
+}
 
 # delete with routing using non-unique lookup vindex
 "delete from multicol_tbl where name = 'foo'"
diff --git a/go/vt/vtgate/planbuilder/testdata/filter_cases.txt b/go/vt/vtgate/planbuilder/testdata/filter_cases.txt
index 754c8a0e1c..685fa720dc 100644
--- a/go/vt/vtgate/planbuilder/testdata/filter_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/filter_cases.txt
@@ -3838,7 +3838,7 @@ Gen4 plan same as above
   "Original": "select * from multicol_tbl where cola = 1",
   "Instructions": {
     "OperatorType": "Route",
-    "Variant": "Equal",
+    "Variant": "SubShard",
     "Keyspace": {
       "Name": "user",
       "Sharded": true
@@ -3875,16 +3875,17 @@ Gen4 plan same as above
   "Original": "select * from multicol_tbl where cola = 1 and colb in (2,3)",
   "Instructions": {
     "OperatorType": "Route",
-    "Variant": "Equal",
+    "Variant": "IN",
     "Keyspace": {
       "Name": "user",
       "Sharded": true
     },
     "FieldQuery": "select * from multicol_tbl where 1 != 1",
-    "Query": "select * from multicol_tbl where cola = 1 and colb in (2, 3)",
+    "Query": "select * from multicol_tbl where cola = 1 and colb in ::__vals1",
     "Table": "multicol_tbl",
     "Values": [
-      "INT64(1)"
+      "INT64(1)",
+      "(INT64(2), INT64(3))"
     ],
     "Vindex": "multicolIdx"
   }