Repository: vitessio/vitess
Pull request #5875 (2020-02-29 14:28:07)

Support owned vindex update if partial column present in set
Based on the thread on slack: https://vitess.slack.com/archives/C0PQY0PTK/p1582885693060100
Propose this change to update lookup column if partial update is in "update set". The remaining vindex column value are taken from owned vindex query result.

Signed-off-by: Harshit Gangal <harshit@planetscale.com>


Patch:
diff --git a/go/vt/vtgate/engine/update.go b/go/vt/vtgate/engine/update.go
index 4fb54b0d84..c5c1c63f48 100644
--- a/go/vt/vtgate/engine/update.go
+++ b/go/vt/vtgate/engine/update.go
@@ -34,12 +34,15 @@ import (
 
 var _ Primitive = (*Update)(nil)
 
+// VindexValues contains changed values for a vindex.
+type VindexValues map[string]sqltypes.PlanValue
+
 // Update represents the instructions to perform an update.
 type Update struct {
 	DML
 
 	// ChangedVindexValues contains values for updated Vindexes during an update statement.
-	ChangedVindexValues map[string][]sqltypes.PlanValue
+	ChangedVindexValues map[string]VindexValues
 
 	// Update does not take inputs
 	noInputs
@@ -60,16 +63,16 @@ func (upd *Update) MarshalJSON() ([]byte, error) {
 	}
 	marshalUpdate := struct {
 		Opcode               string
-		Keyspace             *vindexes.Keyspace              `json:",omitempty"`
-		Query                string                          `json:",omitempty"`
-		Vindex               string                          `json:",omitempty"`
-		Values               []sqltypes.PlanValue            `json:",omitempty"`
-		ChangedVindexValues  map[string][]sqltypes.PlanValue `json:",omitempty"`
-		Table                string                          `json:",omitempty"`
-		OwnedVindexQuery     string                          `json:",omitempty"`
-		KsidVindex           string                          `json:",omitempty"`
-		MultiShardAutocommit bool                            `json:",omitempty"`
-		QueryTimeout         int                             `json:",omitempty"`
+		Keyspace             *vindexes.Keyspace      `json:",omitempty"`
+		Query                string                  `json:",omitempty"`
+		Vindex               string                  `json:",omitempty"`
+		Values               []sqltypes.PlanValue    `json:",omitempty"`
+		ChangedVindexValues  map[string]VindexValues `json:",omitempty"`
+		Table                string                  `json:",omitempty"`
+		OwnedVindexQuery     string                  `json:",omitempty"`
+		KsidVindex           string                  `json:",omitempty"`
+		MultiShardAutocommit bool                    `json:",omitempty"`
+		QueryTimeout         int                     `json:",omitempty"`
 	}{
 		Opcode:               upd.RouteType(),
 		Keyspace:             upd.Keyspace,
@@ -197,29 +200,35 @@ func (upd *Update) updateVindexEntries(vcursor VCursor, bindVars map[string]*que
 		return nil
 	}
 
+	fieldColNumMap := make(map[string]int)
+	for colNum, field := range subQueryResult.Fields {
+		fieldColNumMap[field.Name] = colNum
+	}
+
 	for _, row := range subQueryResult.Rows {
-		colnum := 1 // we start from the first non-vindex col
 		ksid, err := resolveKeyspaceID(vcursor, upd.KsidVindex, row[0])
 		if err != nil {
 			return err
 		}
 		for _, colVindex := range upd.Table.Owned {
-			// Fetch the column values. colnum must keep incrementing.
-			fromIds := make([]sqltypes.Value, 0, len(colVindex.Columns))
-			for range colVindex.Columns {
-				fromIds = append(fromIds, row[colnum])
-				colnum++
-			}
-
 			// Update columns only if they're being changed.
-			if colValues, ok := upd.ChangedVindexValues[colVindex.Name]; ok {
+			if updColValues, ok := upd.ChangedVindexValues[colVindex.Name]; ok {
+				fromIds := make([]sqltypes.Value, 0, len(colVindex.Columns))
 				var vindexColumnKeys []sqltypes.Value
-				for _, colValue := range colValues {
-					resolvedVal, err := colValue.ResolveValue(bindVars)
-					if err != nil {
-						return err
+				for _, vCol := range colVindex.Columns {
+					// Fetch the column values.
+					origColValue := row[fieldColNumMap[vCol.String()]]
+					fromIds = append(fromIds, origColValue)
+					if colValue, exists := updColValues[vCol.String()]; exists {
+						resolvedVal, err := colValue.ResolveValue(bindVars)
+						if err != nil {
+							return err
+						}
+						vindexColumnKeys = append(vindexColumnKeys, resolvedVal)
+					} else {
+						// Set the column value to original as this column in vindex is not updated.
+						vindexColumnKeys = append(vindexColumnKeys, origColValue)
 					}
-					vindexColumnKeys = append(vindexColumnKeys, resolvedVal)
 				}
 
 				if err := colVindex.Vindex.(vindexes.Lookup).Update(vcursor, fromIds, ksid, vindexColumnKeys); err != nil {
diff --git a/go/vt/vtgate/planbuilder/update.go b/go/vt/vtgate/planbuilder/update.go
index fc343014d7..d8c84f4072 100644
--- a/go/vt/vtgate/planbuilder/update.go
+++ b/go/vt/vtgate/planbuilder/update.go
@@ -34,7 +34,7 @@ func buildUpdatePlan(upd *sqlparser.Update, vschema ContextVSchema) (*engine.Upd
 	}
 	eupd := &engine.Update{
 		DML:                 *dml,
-		ChangedVindexValues: make(map[string][]sqltypes.PlanValue),
+		ChangedVindexValues: make(map[string]engine.VindexValues),
 	}
 
 	if dml.Opcode == engine.Unsharded {
@@ -54,10 +54,10 @@ func buildUpdatePlan(upd *sqlparser.Update, vschema ContextVSchema) (*engine.Upd
 // buildChangedVindexesValues adds to the plan all the lookup vindexes that are changing.
 // Updates can only be performed to secondary lookup vindexes with no complex expressions
 // in the set clause.
-func buildChangedVindexesValues(update *sqlparser.Update, colVindexes []*vindexes.ColumnVindex) (map[string][]sqltypes.PlanValue, error) {
-	changedVindexes := make(map[string][]sqltypes.PlanValue)
+func buildChangedVindexesValues(update *sqlparser.Update, colVindexes []*vindexes.ColumnVindex) (map[string]engine.VindexValues, error) {
+	changedVindexes := make(map[string]engine.VindexValues)
 	for i, vindex := range colVindexes {
-		var vindexValues []sqltypes.PlanValue
+		vindexValueMap := make(engine.VindexValues)
 		for _, vcol := range vindex.Columns {
 			// Searching in order of columns in colvindex.
 			found := false
@@ -73,16 +73,13 @@ func buildChangedVindexesValues(update *sqlparser.Update, colVindexes []*vindexe
 				if err != nil {
 					return nil, err
 				}
-				vindexValues = append(vindexValues, pv)
+				vindexValueMap[vcol.String()] = pv
 			}
 		}
-		if len(vindexValues) == 0 {
+		if len(vindexValueMap) == 0 {
 			// Vindex not changing, continue
 			continue
 		}
-		if len(vindexValues) != len(vindex.Columns) {
-			return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: update does not have values for all the columns in vindex (%s)", vindex.Name)
-		}
 
 		if update.Limit != nil && len(update.OrderBy) == 0 {
 			return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: Need to provide order by clause when using limit. Invalid update on vindex: %v", vindex.Name)
@@ -96,7 +93,7 @@ func buildChangedVindexesValues(update *sqlparser.Update, colVindexes []*vindexe
 		if !vindex.Owned {
 			return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: You can only update owned vindexes. Invalid update on vindex: %v", vindex.Name)
 		}
-		changedVindexes[vindex.Name] = vindexValues
+		changedVindexes[vindex.Name] = vindexValueMap
 	}
 
 	return changedVindexes, nil


Test patch:
diff --git a/go/test/endtoend/vtgate/lookup_test.go b/go/test/endtoend/vtgate/lookup_test.go
index 639ef71fef..65c174bf8f 100644
--- a/go/test/endtoend/vtgate/lookup_test.go
+++ b/go/test/endtoend/vtgate/lookup_test.go
@@ -344,6 +344,75 @@ func TestHashLookupMultiInsertIgnore(t *testing.T) {
 	}
 }
 
+func TestConsistentLookupUpdate(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	/* Simple insert. after this dml, the tables will contain the following:
+	t4 (id1, id2):
+	1 2
+	2 2
+	3 3
+	4 3
+
+	t4_id2_idx (id2, id1, keyspace_id:id1):
+	2 1 1
+	2 2 2
+	3 3 3
+	3 4 4
+	*/
+	exec(t, conn, "insert into t4(id1, id2) values(1, '2'), (2, '2'), (3, '3'), (4, '3')")
+	qr := exec(t, conn, "select id1, id2 from t4 order by id1")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT64(1) VARCHAR("2")] [INT64(2) VARCHAR("2")] [INT64(3) VARCHAR("3")] [INT64(4) VARCHAR("3")]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	/* Updating a lookup column. after this dml, the tables will contain the following:
+	t4 (id1, id2):
+	1 42
+	2 2
+	3 3
+	4 3
+
+	t4_id2_idx (id2, id1, keyspace_id:id1):
+	42 1 1
+	2 2 2
+	3 3 3
+	3 4 4
+	*/
+	exec(t, conn, "update t4 set id2 = '42' where id1 = 1")
+	qr = exec(t, conn, "select id1, id2 from t4 order by id1")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT64(1) VARCHAR("42")] [INT64(2) VARCHAR("2")] [INT64(3) VARCHAR("3")] [INT64(4) VARCHAR("3")]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	/* delete one specific keyspace id. after this dml, the tables will contain the following:
+	t4 (id1, id2):
+	2 2
+	3 3
+	4 3
+
+	t4_id2_idx (id2, id1, keyspace_id:id1):
+	2 2 2
+	3 3 3
+	3 4 4
+	*/
+	exec(t, conn, "delete from t4 where id2 = '42'")
+	qr = exec(t, conn, "select * from t4 where id2 = '42'")
+	require.Empty(t, qr.Rows)
+	qr = exec(t, conn, "select * from t4_id2_idx where id2 = '42'")
+	require.Empty(t, qr.Rows)
+
+	// delete all the rows.
+	exec(t, conn, "delete from t4")
+	qr = exec(t, conn, "select * from t4")
+	require.Empty(t, qr.Rows)
+	qr = exec(t, conn, "select * from t4_id2_idx")
+	require.Empty(t, qr.Rows)
+}
+
 func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
 	t.Helper()
 	qr, err := conn.ExecuteFetch(query, 1000, true)
diff --git a/go/test/endtoend/vtgate/main_test.go b/go/test/endtoend/vtgate/main_test.go
index 1903fdfcb6..5ca001fb8a 100644
--- a/go/test/endtoend/vtgate/main_test.go
+++ b/go/test/endtoend/vtgate/main_test.go
@@ -81,12 +81,28 @@ create table t3_id7_idx(
 	id7 bigint,
 	id6 bigint,
     primary key(id)
+) Engine=InnoDB;
+
+create table t4(
+	id1 bigint,
+	id2 varchar(10),
+	primary key(id1)
+) Engine=InnoDB;
+
+create table t4_id2_idx(
+	id2 varchar(10),
+	id1 bigint,
+	keyspace_id varbinary(50),
+    primary key(id2, id1)
 ) Engine=InnoDB;`
 
 	VSchema = `
 {
   "sharded": true,
   "vindexes": {
+    "unicode_loose_md5" : {
+	  "type": "unicode_loose_md5"
+    },
     "hash": {
       "type": "hash"
     },
@@ -117,6 +133,15 @@ create table t3_id7_idx(
         "to": "id6"
       },
       "owner": "t3"
+    },
+    "t4_id2_vdx": {
+      "type": "consistent_lookup",
+      "params": {
+        "table": "t4_id2_idx",
+        "from": "id2,id1",
+        "to": "keyspace_id"
+      },
+      "owner": "t4"
     }
   },
   "tables": {
@@ -179,6 +204,26 @@ create table t3_id7_idx(
           "name": "hash"
         }
       ]
+    },
+	"t4": {
+      "column_vindexes": [
+        {
+          "column": "id1",
+          "name": "hash"
+        },
+        {
+          "columns": ["id2", "id1"],
+          "name": "t4_id2_vdx"
+        }
+      ]
+    },
+    "t4_id2_idx": {
+      "column_vindexes": [
+        {
+          "column": "id2",
+          "name": "unicode_loose_md5"
+        }
+      ]
     },
     "vstream_test": {
       "column_vindexes": [
diff --git a/go/vt/vtgate/engine/update_test.go b/go/vt/vtgate/engine/update_test.go
index b73590a7ca..356c9c1074 100644
--- a/go/vt/vtgate/engine/update_test.go
+++ b/go/vt/vtgate/engine/update_test.go
@@ -203,15 +203,14 @@ func TestUpdateEqualChangedVindex(t *testing.T) {
 			OwnedVindexQuery: "dummy_subquery",
 			KsidVindex:       ks.Vindexes["hash"].(vindexes.SingleColumn),
 		},
-		ChangedVindexValues: map[string][]sqltypes.PlanValue{
-			"twocol": {{
-				Value: sqltypes.NewInt64(1),
-			}, {
-				Value: sqltypes.NewInt64(2),
-			}},
-			"onecol": {{
-				Value: sqltypes.NewInt64(3),
-			}},
+		ChangedVindexValues: map[string]VindexValues{
+			"twocol": {
+				"c1": {Value: sqltypes.NewInt64(1)},
+				"c2": {Value: sqltypes.NewInt64(2)},
+			},
+			"onecol": {
+				"c3": {Value: sqltypes.NewInt64(3)},
+			},
 		},
 	}
 
@@ -309,15 +308,14 @@ func TestUpdateScatterChangedVindex(t *testing.T) {
 			OwnedVindexQuery: "dummy_subquery",
 			KsidVindex:       ks.Vindexes["hash"].(vindexes.SingleColumn),
 		},
-		ChangedVindexValues: map[string][]sqltypes.PlanValue{
-			"twocol": {{
-				Value: sqltypes.NewInt64(1),
-			}, {
-				Value: sqltypes.NewInt64(2),
-			}},
-			"onecol": {{
-				Value: sqltypes.NewInt64(3),
-			}},
+		ChangedVindexValues: map[string]VindexValues{
+			"twocol": {
+				"c1": {Value: sqltypes.NewInt64(1)},
+				"c2": {Value: sqltypes.NewInt64(2)},
+			},
+			"onecol": {
+				"c3": {Value: sqltypes.NewInt64(3)},
+			},
 		},
 	}
 
diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.txt b/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
index 1de0c561b2..4c84422be4 100644
--- a/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/dml_cases.txt
@@ -229,9 +229,9 @@
       1
     ],
     "ChangedVindexValues": {
-      "email_user_map": [
-        "juan@vitess.io"
-      ]
+      "email_user_map": {
+        "email": "juan@vitess.io"
+      }
     },
     "Table": "user_metadata",
     "OwnedVindexQuery": "select user_id, email, address from user_metadata where user_id = 1 for update",
@@ -259,12 +259,12 @@
       1
     ],
     "ChangedVindexValues": {
-      "address_user_map": [
-        "155 5th street"
-      ],
-      "email_user_map": [
-        "juan@vitess.io"
-      ]
+      "address_user_map": {
+        "address": "155 5th street"
+      },
+      "email_user_map": {
+        "email": "juan@vitess.io"
+      }
     },
     "Table": "user_metadata",
     "OwnedVindexQuery": "select user_id, email, address from user_metadata where user_id = 1 for update",
@@ -288,9 +288,9 @@
       1
     ],
     "ChangedVindexValues": {
-      "email_user_map": [
-        "juan@vitess.io"
-      ]
+      "email_user_map": {
+        "email": "juan@vitess.io"
+      }
     },
     "Table": "user_metadata",
     "OwnedVindexQuery": "select user_id, email, address from user_metadata where user_id = 1 order by user_id asc limit 10 for update",
@@ -1777,10 +1777,10 @@
       1
     ],
     "ChangedVindexValues": {
-      "colb_colc_map": [
-        1,
-        2
-      ]
+      "colb_colc_map": {
+        "column_b": 1,
+        "column_c": 2
+      }
     },
     "Table": "multicolvin",
     "OwnedVindexQuery": "select kid, column_a, column_b, column_c from multicolvin where kid = 1 for update",
@@ -1804,13 +1804,13 @@
       1
     ],
     "ChangedVindexValues": {
-      "cola_map": [
-        0
-      ],
-      "colb_colc_map": [
-        1,
-        2
-      ]
+      "cola_map": {
+        "column_a": 0
+      },
+      "colb_colc_map": {
+        "column_b": 1,
+        "column_c": 2
+      }
     },
     "Table": "multicolvin",
     "OwnedVindexQuery": "select kid, column_a, column_b, column_c from multicolvin where kid = 1 for update",
@@ -2106,9 +2106,9 @@
       1
     ],
     "ChangedVindexValues": {
-      "name_user_map": [
-        null
-      ]
+      "name_user_map": {
+        "Name": null
+      }
     },
     "Table": "user",
     "OwnedVindexQuery": "select Id, Name, Costly from user where id = 1 for update",
@@ -2143,9 +2143,9 @@
     },
     "Query": "update user set name = null where id in (1, 2, 3)",
     "ChangedVindexValues": {
-      "name_user_map": [
-        null
-      ]
+      "name_user_map": {
+        "Name": null
+      }
     },
     "Table": "user",
     "OwnedVindexQuery": "select Id, Name, Costly from user where id in (1, 2, 3) for update",
@@ -2165,9 +2165,9 @@
     },
     "Query": "update user set name = null",
     "ChangedVindexValues": {
-      "name_user_map": [
-        null
-      ]
+      "name_user_map": {
+        "Name": null
+      }
     },
     "Table": "user",
     "OwnedVindexQuery": "select Id, Name, Costly from user for update",
@@ -2187,9 +2187,9 @@
     },
     "Query": "update user set name = null where id + 1 = 2",
     "ChangedVindexValues": {
-      "name_user_map": [
-        null
-      ]
+      "name_user_map": {
+        "Name": null
+      }
     },
     "Table": "user",
     "OwnedVindexQuery": "select Id, Name, Costly from user where id + 1 = 2 for update",
@@ -2300,3 +2300,29 @@
     "KsidVindex": "user_index"
   }
 }
+
+# update multi column vindex, without values for all the vindex columns
+"update multicolvin set column_c = 2 where kid = 1"
+{
+  "Original": "update multicolvin set column_c = 2 where kid = 1",
+  "Instructions": {
+    "Opcode": "UpdateEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update multicolvin set column_c = 2 where kid = 1",
+    "Vindex": "kid_index",
+    "Values": [
+      1
+    ],
+    "ChangedVindexValues": {
+      "colb_colc_map": {
+        "column_c": 2
+      }
+    },
+    "Table": "multicolvin",
+    "OwnedVindexQuery": "select kid, column_a, column_b, column_c from multicolvin where kid = 1 for update",
+    "KsidVindex": "kid_index"
+  }
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
index 0bf10af7c8..7bccf1373f 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
@@ -259,10 +259,6 @@
 "update user_metadata set email = 'juan@vitess.io' where user_id = 1 limit 10"
 "unsupported: Need to provide order by clause when using limit. Invalid update on vindex: email_user_map"
 
-# update multi column vindex, without values for all the vindex columns
-"update multicolvin set column_c = 2 where kid = 1"
-"unsupported: update does not have values for all the columns in vindex (colb_colc_map)"
-
 # cross-shard update tables
 "update (select id from user) as u set id = 4"
 "unsupported: subqueries in sharded DML"