release-20.2: sql: fix UPSERT with columns after partial index PUT and DEL columns

pkg/sql/delete.go

@@ -158,9 +158,9 @@ func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	// satisfy the predicate and therefore do not exist in the partial index.
 	// This set is passed as a argument to tableDeleter.row below.
 	var pm row.PartialIndexUpdateHelper
-	partialIndexOrds := d.run.td.tableDesc().PartialIndexOrds()
-	if !partialIndexOrds.Empty() {
-		partialIndexDelVals := sourceVals[d.run.partialIndexDelValsOffset:]
+	if n := d.run.td.tableDesc().PartialIndexOrds().Len(); n > 0 {
+		offset := d.run.partialIndexDelValsOffset
+		partialIndexDelVals := sourceVals[offset : offset+n]
 
 		err := pm.Init(tree.Datums{}, partialIndexDelVals, d.run.td.tableDesc())
 		if err != nil {

pkg/sql/insert.go

@@ -130,9 +130,9 @@ func (r *insertRun) processSourceRow(params runParams, rowVals tree.Datums) erro
 	// written to when they are partial indexes and the row does not satisfy the
 	// predicate. This set is passed as a parameter to tableInserter.row below.
 	var pm row.PartialIndexUpdateHelper
-	partialIndexOrds := r.ti.tableDesc().PartialIndexOrds()
-	if !partialIndexOrds.Empty() {
-		partialIndexPutVals := rowVals[len(r.insertCols)+r.checkOrds.Len():]
+	if n := r.ti.tableDesc().PartialIndexOrds().Len(); n > 0 {
+		offset := len(r.insertCols) + r.checkOrds.Len()
+		partialIndexPutVals := rowVals[offset : offset+n]
 
 		err := pm.Init(partialIndexPutVals, tree.Datums{}, r.ti.tableDesc())
 		if err != nil {

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -1428,3 +1428,31 @@ CREATE TABLE t58390 (
 
 statement ok
 ALTER TABLE t58390 ALTER PRIMARY KEY USING COLUMNS (b, a)
+
+# Regression tests for #61414. Upsert execution should not error if partial
+# index PUT and DEL columns are not the last columns in the input of the
+# mutation.
+
+statement ok
+create table t61414_a (
+  k INT PRIMARY KEY
+)
+
+statement ok
+create table t61414_b (
+  k INT PRIMARY KEY,
+  a STRING,
+  b INT REFERENCES t61414_a(k),
+  UNIQUE INDEX (a, b),
+  INDEX (a) WHERE a = 'x'
+)
+
+statement ok
+INSERT INTO t61414_a VALUES (2)
+
+statement ok
+INSERT INTO t61414_b (k, a, b)
+VALUES (1, 'a', 2)
+ON CONFLICT (a, b) DO UPDATE SET a = excluded.a
+WHERE t61414_b.a = 'x'
+RETURNING k

pkg/sql/update.go

@@ -295,12 +295,11 @@ func (u *updateNode) processSourceRow(params runParams, sourceVals tree.Datums)
 	// Create a set of partial index IDs to not add entries or remove entries
 	// from.
 	var pm row.PartialIndexUpdateHelper
-	partialIndexOrds := u.run.tu.tableDesc().PartialIndexOrds()
-	if !partialIndexOrds.Empty() {
-		partialIndexValOffset := len(u.run.tu.ru.FetchCols) + len(u.run.tu.ru.UpdateCols) + u.run.checkOrds.Len() + u.run.numPassthrough
-		partialIndexVals := sourceVals[partialIndexValOffset:]
-		partialIndexPutVals := partialIndexVals[:partialIndexOrds.Len()]
-		partialIndexDelVals := partialIndexVals[partialIndexOrds.Len() : partialIndexOrds.Len()*2]
+	if n := u.run.tu.tableDesc().PartialIndexOrds().Len(); n > 0 {
+		offset := len(u.run.tu.ru.FetchCols) + len(u.run.tu.ru.UpdateCols) + u.run.checkOrds.Len() + u.run.numPassthrough
+		partialIndexVals := sourceVals[offset:]
+		partialIndexPutVals := partialIndexVals[:n]
+		partialIndexDelVals := partialIndexVals[n : n*2]
 
 		err := pm.Init(partialIndexPutVals, partialIndexDelVals, u.run.tu.tableDesc())
 		if err != nil {

pkg/sql/upsert.go

@@ -146,15 +146,14 @@ func (n *upsertNode) processSourceRow(params runParams, rowVals tree.Datums) err
 
 	// Create a set of partial index IDs to not add or remove entries from.
 	var pm row.PartialIndexUpdateHelper
-	partialIndexOrds := n.run.tw.tableDesc().PartialIndexOrds()
-	if !partialIndexOrds.Empty() {
-		partialIndexValOffset := len(n.run.insertCols) + len(n.run.tw.fetchCols) + len(n.run.tw.updateCols) + n.run.checkOrds.Len()
+	if numPartialIndexes := n.run.tw.tableDesc().PartialIndexOrds().Len(); numPartialIndexes > 0 {
+		offset := len(n.run.insertCols) + len(n.run.tw.fetchCols) + len(n.run.tw.updateCols) + n.run.checkOrds.Len()
 		if n.run.tw.canaryOrdinal != -1 {
-			partialIndexValOffset++
+			offset++
 		}
-		partialIndexVals := rowVals[partialIndexValOffset:]
-		partialIndexPutVals := partialIndexVals[:len(partialIndexVals)/2]
-		partialIndexDelVals := partialIndexVals[len(partialIndexVals)/2:]
+		partialIndexVals := rowVals[offset:]
+		partialIndexPutVals := partialIndexVals[:numPartialIndexes]
+		partialIndexDelVals := partialIndexVals[numPartialIndexes : numPartialIndexes*2]
 
 		err := pm.Init(partialIndexPutVals, partialIndexDelVals, n.run.tw.tableDesc())
 		if err != nil {
@@ -163,7 +162,7 @@ func (n *upsertNode) processSourceRow(params runParams, rowVals tree.Datums) err
 
 		// Truncate rowVals so that it no longer includes partial index predicate
 		// values.
-		rowVals = rowVals[:partialIndexValOffset]
+		rowVals = rowVals[:offset]
 	}
 
 	// Verify the CHECK constraints by inspecting boolean columns from the input that