row, sql: implement mutations on vector indexes

This patch plumbs the output of the vector search operators into rowenc
for encoding into vector indexes. The output of these vector search
operators is included in the row values for mutation operators after
partial index values and are plumbed into a new vector index update
helper, which tracks the column values until they're needed by rowenc.

While we're here, we homogonize how
pkg/sql/{delete,insert,update,upsert}.go consume row values, hopefully
improving legibility.

Also fixed are a couple of bugs in the vector search and row encoding.

Epic: CRDB-42943
Release note: None

Author: mw5h

pkg/crosscluster/logical/lww_kv_processor.go

@@ -476,13 +476,15 @@ func (p *kvTableWriter) insertRow(ctx context.Context, b *kv.Batch, after cdceve
 
 	var ph row.PartialIndexUpdateHelper
 	// TODO(dt): support partial indexes.
+	var vh row.VectorIndexUpdateHelper
+	// TODO(mw5h, drewk): support vector indexes.
 	oth := &row.OriginTimestampCPutHelper{
 		OriginTimestamp: after.MvccTimestamp,
 		// TODO(ssd): We should choose this based by comparing the cluster IDs of the source
 		// and destination clusters.
 		// ShouldWinTie: true,
 	}
-	return p.ri.InsertRow(ctx, &row.KVBatchAdapter{Batch: b}, p.newVals, ph, oth, row.CPutOp, false /* traceKV */)
+	return p.ri.InsertRow(ctx, &row.KVBatchAdapter{Batch: b}, p.newVals, ph, vh, oth, row.CPutOp, false /* traceKV */)
 }
 
 func (p *kvTableWriter) updateRow(
@@ -497,13 +499,15 @@ func (p *kvTableWriter) updateRow(
 
 	var ph row.PartialIndexUpdateHelper
 	// TODO(dt): support partial indexes.
+	var vh row.VectorIndexUpdateHelper
+	// TODO(mw5h, drewk): support vector indexes.
 	oth := &row.OriginTimestampCPutHelper{
 		OriginTimestamp: after.MvccTimestamp,
 		// TODO(ssd): We should choose this based by comparing the cluster IDs of the source
 		// and destination clusters.
 		// ShouldWinTie: true,
 	}
-	_, err := p.ru.UpdateRow(ctx, b, p.oldVals, p.newVals, ph, oth, false)
+	_, err := p.ru.UpdateRow(ctx, b, p.oldVals, p.newVals, ph, vh, oth, false)
 	return err
 }
 
@@ -516,6 +520,8 @@ func (p *kvTableWriter) deleteRow(
 
 	var ph row.PartialIndexUpdateHelper
 	// TODO(dt): support partial indexes.
+	var vh row.VectorIndexUpdateHelper
+	// TODO(mw5h, drewk): support vector indexes.
 	oth := &row.OriginTimestampCPutHelper{
 		PreviousWasDeleted: before.IsDeleted(),
 		OriginTimestamp:    after.MvccTimestamp,
@@ -524,7 +530,7 @@ func (p *kvTableWriter) deleteRow(
 		// ShouldWinTie: true,
 	}
 
-	return p.rd.DeleteRow(ctx, b, p.oldVals, ph, oth, false)
+	return p.rd.DeleteRow(ctx, b, p.oldVals, ph, vh, oth, false)
 }
 
 func (p *kvTableWriter) fillOld(vals cdcevent.Row) error {

pkg/sql/backfill/backfill.go

@@ -398,12 +398,13 @@ func (cb *ColumnBackfiller) RunColumnBackfillChunk(
 		}
 		copy(oldValues, fetchedValues)
 
-		// No existing secondary indexes will be updated by adding or dropping a
-		// column. It is safe to use an empty PartialIndexUpdateHelper in this
-		// case.
+		// No existing secondary indexes will be updated by adding or dropping a column.
+		// It is safe to use an empty PartialIndexUpdateHelper and
+		// VectorIndexUpdateHelper in this case.
 		var pm row.PartialIndexUpdateHelper
+		var vh row.VectorIndexUpdateHelper
 		if _, err := ru.UpdateRow(
-			ctx, b, oldValues, updateValues, pm, nil, traceKV,
+			ctx, b, oldValues, updateValues, pm, vh, nil, traceKV,
 		); err != nil {
 			return roachpb.Key{}, err
 		}

pkg/sql/colenc/encode.go

@@ -432,6 +432,11 @@ func (b *BatchEncoder) encodeSecondaryIndex(ctx context.Context, ind catalog.Ind
 		// Since the inverted indexes generate multiple keys per row just handle them
 		// separately.
 		return b.encodeInvertedSecondaryIndex(ctx, ind, kys, b.extraKeys)
+	} else if ind.GetType() == idxtype.VECTOR {
+		// TODO(mw5h, drewk): Implement encoding for vector indexes. This is complicated
+		// by vector indexes needing output from the mutation search operator, which will
+		// need to be plumbed through to the encoder.
+		return errors.AssertionFailedf("vector indexes not supported")
 	} else {
 		keyAndSuffixCols := b.rh.TableDesc.IndexFetchSpecKeyAndSuffixColumns(ind)
 		keyCols := keyAndSuffixCols[:ind.NumKeyColumns()]

pkg/sql/colenc/encode_test.go

@@ -606,8 +606,9 @@ func buildRowKVs(
 	}
 	p := &capturePutter{}
 	var pm row.PartialIndexUpdateHelper
+	var vh row.VectorIndexUpdateHelper
 	for _, d := range datums {
-		if err := inserter.InsertRow(context.Background(), p, d, pm, nil, row.CPutOp, true /* traceKV */); err != nil {
+		if err := inserter.InsertRow(context.Background(), p, d, pm, vh, nil, row.CPutOp, true /* traceKV */); err != nil {
 			return kvs{}, err
 		}
 	}

pkg/sql/copy_from.go

@@ -418,6 +418,10 @@ func (c *copyMachine) canSupportVectorized(table catalog.TableDescriptor) bool {
 	if c.p.SessionData().VectorizeMode == sessiondatapb.VectorizeOff {
 		return false
 	}
+	// Columnar vector index encoding is not yet supported.
+	if len(table.VectorIndexes()) > 0 {
+		return false
+	}
 	// Vectorized COPY doesn't support foreign key checks, no reason it couldn't
 	// but it doesn't work right now because we don't have the ability to
 	// hold the results in a bufferNode. We wouldn't want to enable it

pkg/sql/create_table.go

@@ -612,11 +612,12 @@ func (n *createTableNode) startExec(params runParams) error {
 				// Populate the buffer.
 				copy(rowBuffer, n.input.Values())
 
-				// CREATE TABLE AS does not copy indexes from the input table.
-				// An empty row.PartialIndexUpdateHelper is used here because
-				// there are no indexes, partial or otherwise, to update.
+				// CREATE TABLE AS does not copy indexes from the input table. Empty
+				// partial and vector index helpers are used here because there are no
+				// indexes, partial, vector, or otherwise, to update.
 				var pm row.PartialIndexUpdateHelper
-				if err := ti.row(params.ctx, rowBuffer, pm, params.extendedEvalCtx.Tracing.KVTracingEnabled()); err != nil {
+				var vh row.VectorIndexUpdateHelper
+				if err := ti.row(params.ctx, rowBuffer, pm, vh, params.extendedEvalCtx.Tracing.KVTracingEnabled()); err != nil {
 					return err
 				}
 			}

pkg/sql/delete.go

@@ -146,29 +146,33 @@ func (d *deleteNode) BatchedNext(params runParams) (bool, error) {
 // processSourceRow processes one row from the source for deletion and, if
 // result rows are needed, saves it in the result row container
 func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums) error {
+	// Remove extra columns for partial index predicate values and AFTER triggers.
+	deleteVals := sourceVals[:len(d.run.td.rd.FetchCols)+d.run.numPassthrough]
+	sourceVals = sourceVals[len(deleteVals):]
+
 	// Create a set of partial index IDs to not delete from. Indexes should not
 	// be deleted from when they are partial indexes and the row does not
 	// 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
-	deleteCols := len(d.run.td.rd.FetchCols) + d.run.numPassthrough
 	if n := len(d.run.td.tableDesc().PartialIndexes()); n > 0 {
-		partialIndexDelVals := sourceVals[deleteCols : deleteCols+n]
-
-		err := pm.Init(nil /* partialIndexPutVals */, partialIndexDelVals, d.run.td.tableDesc())
+		err := pm.Init(nil /* partialIndexPutVals */, sourceVals[:n], d.run.td.tableDesc())
 		if err != nil {
 			return err
 		}
+		sourceVals = sourceVals[n:]
 	}
 
-	if len(sourceVals) > deleteCols {
-		// Remove extra columns for partial index predicate values and AFTER
-		// triggers.
-		sourceVals = sourceVals[:deleteCols]
+	// Keep track of the vector index partitions to update. This information is
+	// passed to tableInserter.row below.
+	var vh row.VectorIndexUpdateHelper
+	if n := len(d.run.td.tableDesc().VectorIndexes()); n > 0 {
+		vh.InitForDel(sourceVals[:n], d.run.td.tableDesc())
+		sourceVals = sourceVals[n:]
 	}
 
 	// Queue the deletion in the KV batch.
-	if err := d.run.td.row(params.ctx, sourceVals, pm, d.run.traceKV); err != nil {
+	if err := d.run.td.row(params.ctx, deleteVals, pm, vh, d.run.traceKV); err != nil {
 		return err
 	}
 
@@ -188,7 +192,7 @@ func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums)
 				if retIdx >= largestRetIdx {
 					largestRetIdx = retIdx
 				}
-				resultValues[retIdx] = sourceVals[i]
+				resultValues[retIdx] = deleteVals[i]
 			}
 		}
 
@@ -198,7 +202,7 @@ func (d *deleteNode) processSourceRow(params runParams, sourceVals tree.Datums)
 		if d.run.numPassthrough > 0 {
 			passthroughBegin := len(d.run.td.rd.FetchCols)
 			passthroughEnd := passthroughBegin + d.run.numPassthrough
-			passthroughValues := sourceVals[passthroughBegin:passthroughEnd]
+			passthroughValues := deleteVals[passthroughBegin:passthroughEnd]
 
 			for i := 0; i < d.run.numPassthrough; i++ {
 				largestRetIdx++

pkg/sql/distsql_physical_planner.go

@@ -4622,7 +4622,8 @@ func (dsp *DistSQLPlanner) planVectorMutationSearch(
 		suffixKeyColumnOrdinals[i] = uint32(p.PlanToStreamColMap[col])
 	}
 	keyAndSuffixCols := planInfo.table.IndexFetchSpecKeyAndSuffixColumns(planInfo.index)
-	prefixKeyCols := keyAndSuffixCols[planInfo.index.NumKeyColumns()-1:]
+	// Prefix cols do not include the vector column, which is the last key column.
+	prefixKeyCols := keyAndSuffixCols[:planInfo.index.NumKeyColumns()-1]
 	suffixKeyCols := keyAndSuffixCols[planInfo.index.NumKeyColumns():]
 	spec := &execinfrapb.VectorMutationSearchSpec{
 		PrefixKeyColumnOrdinals:  prefixKeyColumnOrdinals,

pkg/sql/insert.go

@@ -202,31 +202,39 @@ func (r *insertRun) processSourceRow(params runParams, rowVals tree.Datums) erro
 		return err
 	}
 
+	rowVals = rowVals[len(insertVals):]
+
+	// Verify the CHECK constraint results, if any.
+	if n := r.checkOrds.Len(); n > 0 {
+		if err := checkMutationInput(
+			params.ctx, params.p.EvalContext(), &params.p.semaCtx, params.p.SessionData(),
+			r.ti.tableDesc(), r.checkOrds, rowVals[:n],
+		); err != nil {
+			return err
+		}
+		rowVals = rowVals[n:]
+	}
+
 	// Create a set of partial index IDs to not write to. Indexes should not be
 	// 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
 	if n := len(r.ti.tableDesc().PartialIndexes()); n > 0 {
-		offset := len(r.insertCols) + r.checkOrds.Len()
-		partialIndexPutVals := rowVals[offset : offset+n]
-
-		err := pm.Init(partialIndexPutVals, nil /* partialIndexDelVals */, r.ti.tableDesc())
+		err := pm.Init(rowVals[:n], nil /* partialIndexDelVals */, r.ti.tableDesc())
 		if err != nil {
 			return err
 		}
+		rowVals = rowVals[n:]
 	}
 
-	// Verify the CHECK constraint results, if any.
-	if n := r.checkOrds.Len(); n > 0 {
-		// CHECK constraint results are after the insert columns.
-		offset := len(r.insertCols)
-		checkVals := rowVals[offset : offset+n]
-		if err := checkMutationInput(
-			params.ctx, params.p.EvalContext(), &params.p.semaCtx, params.p.SessionData(),
-			r.ti.tableDesc(), r.checkOrds, checkVals,
-		); err != nil {
-			return err
-		}
+	// Keep track of the vector index partitions to update, as well as the
+	// quantized vectors. This information is passed to tableInserter.row below.
+	// Input is one partition key per vector index followed by one quantized vector
+	// per index.
+	var vh row.VectorIndexUpdateHelper
+	if n := len(r.ti.tableDesc().VectorIndexes()); n > 0 {
+		vh.InitForPut(rowVals[:n], rowVals[n:n*2], r.ti.tableDesc())
+		rowVals = rowVals[n*2:]
 	}
 
 	// Error out the insert if the enforce_home_region session setting is on and
@@ -236,7 +244,7 @@ func (r *insertRun) processSourceRow(params runParams, rowVals tree.Datums) erro
 	}
 
 	// Queue the insert in the KV batch.
-	if err := r.ti.row(params.ctx, insertVals, pm, r.traceKV); err != nil {
+	if err := r.ti.row(params.ctx, insertVals, pm, vh, r.traceKV); err != nil {
 		return err
 	}
 

pkg/sql/logictest/logic.go

@@ -319,16 +319,17 @@ import (
 //            asynchronously, subsequent queries that depend on the state of
 //            the query should be run with the "retry" option to ensure
 //            deterministic test results.
-//      - kvtrace: runs the query and compares against the results of the
-//            kv operations trace of the query. kvtrace optionally accepts
-//            arguments of the form kvtrace(op,op,...). Op is one of
-//            the accepted k/v arguments such as 'CPut', 'Scan' etc. It
-//            also accepts arguments of the form 'prefix=...'. For example,
-//            if kvtrace(CPut,Del,prefix=/Table/54,prefix=/Table/55), the
-//            results will be filtered to contain messages starting with
-//            CPut /Table/54, CPut /Table/55, Del /Table/54, Del /Table/55.
-//            Tenant IDs do not need to be included in prefixes and will be
-//            removed from results. Cannot be combined with noticetrace.
+//      - kvtrace: runs the query and compares against the results of the kv
+//            operations trace of the query. kvtrace optionally accepts arguments of the
+//            form kvtrace(op,op,...). Op is one of the accepted k/v arguments such as
+//            'CPut', 'Scan' etc. It also accepts arguments of the form 'prefix=...'. For
+//            example, if kvtrace(CPut,Del,prefix=/Table/54,prefix=/Table/55), the results
+//            will be filtered to contain messages starting with CPut /Table/54, CPut
+//            /Table/55, Del /Table/54, Del /Table/55. The 'redactbytes' will redact the
+//            contents of /BYTES/ values to prevent test flakiness in the presence of
+//            nondeterminism and/or processor architecture differences. Tenant IDs do not
+//            need to be included in prefixes and will be removed from results. Cannot be
+//            combined with noticetrace.
 //      - noticetrace: runs the query and compares only the notices that
 //						appear. Cannot be combined with kvtrace.
 //      - nodeidx=N: runs the query on node N of the cluster.
@@ -944,6 +945,10 @@ type logicQuery struct {
 	// the particular operation types to filter on, such as CPut or Del.
 	kvOpTypes        []string
 	keyPrefixFilters []string
+	// kvtraceRedactBytes can only be used when kvtrace is true. When active, BYTES
+	// values in keys are expunged from output (to prevent test failures where the
+	// BYTES value is nondeterministic or architecture dependent).
+	kvtraceRedactBytes bool
 
 	// nodeIdx determines which node on the cluster to execute a query on for the given query.
 	nodeIdx int
@@ -2835,6 +2840,8 @@ func (t *logicTest) processSubtest(
 										matched = "/Tenant/%" + matched
 									}
 									query.keyPrefixFilters = append(query.keyPrefixFilters, matched)
+								} else if c == "redactbytes" {
+									query.kvtraceRedactBytes = true
 								} else if isAllowedKVOp(c) {
 									query.kvOpTypes = append(query.kvOpTypes, c)
 								} else {
@@ -2872,6 +2879,7 @@ func (t *logicTest) processSubtest(
 							query.kvtrace = true
 							query.kvOpTypes = nil
 							query.keyPrefixFilters = nil
+							query.kvtraceRedactBytes = false
 
 						case "noticetrace":
 							query.noticetrace = true
@@ -3017,7 +3025,13 @@ func (t *logicTest) processSubtest(
 
 					projection := `message`
 					if len(t.tenantApps) != 0 || t.cluster.StartedDefaultTestTenant() {
-						projection = `regexp_replace(message, '/Tenant/\d+', '')`
+						projection = `regexp_replace(message, '/Tenant/\d+', '', 'g')`
+					}
+					if query.kvtraceRedactBytes {
+						projection = fmt.Sprintf(
+							`regexp_replace(%s, '/BYTES/0x[abcdef\d]+', '/BYTES/:redacted:', 'g')`,
+							projection,
+						)
 					}
 					queryPrefix := fmt.Sprintf(`SELECT %s FROM [SHOW KV TRACE FOR SESSION] `, projection)
 					buildQuery := func(ops []string, keyFilters []string) string {