feat: make gen4 plan subsharding queries better (#10151)

* feat: make gen4 plan subsharding queries better

Signed-off-by: Andres Taylor <andres@planetscale.com>

* chore: updated tests and engine primitives to handle Subshard opcode

Signed-off-by: Andres Taylor <andres@planetscale.com>

* test: update tests after adding Subshard

Signed-off-by: Andres Taylor <andres@planetscale.com>

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.

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,

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)

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)

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.

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,

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)

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
 	}

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{}

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'"

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"
   }

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)
 			}