vecindex: clone treeKey when creating fixups

Previously, the C-SPANN index assumed that it owned the memory for
the vector index prefix value (called the "tree key" in C-SPANN).
However, the backfiller assumed the opposite, and reuses the memory
between SearchForInsert calls. This commit resolves the conflict by
always cloning the tree key when it's enqueued as part of a fixup,
so that callers can reuse the memory.

Fixes: #145261
Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/backfill/backfill_test.go

@@ -181,3 +181,63 @@ func TestConcurrentOperationsDuringVectorIndexCreation(t *testing.T) {
 	sqlDB.QueryRow(t, `SELECT id FROM vectors@vec_idx ORDER BY vec <-> '[1, 2, 3]' LIMIT 1`).Scan(&id)
 	require.Equal(t, 1, id)
 }
+
+// Regression for issue #145261: vector index backfill with a prefix column
+// crashes the node.
+func TestVectorIndexWithPrefixBackfill(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	srv, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	sqlDB := sqlutils.MakeSQLRunner(db)
+
+	// Enable vector indexes.
+	sqlDB.Exec(t, `SET CLUSTER SETTING feature.vector_index.enabled = true`)
+
+	// Create a table with a vector column + a prefix column.
+	sqlDB.Exec(t, `
+		CREATE TABLE items (
+			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+			customer_id INT NOT NULL,
+			name TEXT,
+			embedding VECTOR(3)
+		);
+	`)
+
+	// Generate 10 customer id's, each with 100 vectors.
+	sqlDB.Exec(t, `
+		INSERT INTO items (customer_id, name, embedding)
+		SELECT
+			(i % 10) + 1 AS customer_id,
+			'Item ' || i,
+			ARRAY[random(), random(), random()]::vector
+		FROM generate_series(1, 1000) AS s(i);
+	`)
+
+	// Create the vector index with a small partition size so that the trees
+	// have more levels and splits to get there.
+	sqlDB.Exec(t, `
+		CREATE VECTOR INDEX ON items (customer_id, embedding)
+		WITH (min_partition_size=2, max_partition_size=8, build_beam_size=2);
+	`)
+
+	// Ensure that each customer has 100 vectors.
+	for i := range 10 {
+		func() {
+			rows := sqlDB.Query(t, `
+				SELECT id FROM items
+				WHERE customer_id = $1
+				ORDER BY embedding <-> $2
+				LIMIT 200`, i+1, "[0, 0, 0]")
+			defer rows.Close()
+
+			count := 0
+			for rows.Next() {
+				count++
+			}
+			require.Equal(t, 100, count)
+		}()
+	}
+}

pkg/sql/vecindex/cspann/fixup_processor.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"runtime"
+	"slices"
 	"sync"
 	"time"
 
@@ -290,7 +291,8 @@ func (fp *FixupProcessor) AddDeleteVector(
 	ctx context.Context, treeKey TreeKey, partitionKey PartitionKey, vectorKey KeyBytes,
 ) {
 	fp.addFixup(ctx, fixup{
-		TreeKey:      treeKey,
+		// Clone the tree key, since we don't own the memory.
+		TreeKey:      slices.Clone(treeKey),
 		Type:         vectorDeleteFixup,
 		PartitionKey: partitionKey,
 		VectorKey:    vectorKey,
@@ -306,7 +308,8 @@ func (fp *FixupProcessor) AddSplit(
 	singleStep bool,
 ) {
 	fp.addFixup(ctx, fixup{
-		TreeKey:            treeKey,
+		// Clone the tree key, since we don't own the memory.
+		TreeKey:            slices.Clone(treeKey),
 		Type:               splitFixup,
 		ParentPartitionKey: parentPartitionKey,
 		PartitionKey:       partitionKey,
@@ -323,7 +326,8 @@ func (fp *FixupProcessor) AddMerge(
 	singleStep bool,
 ) {
 	fp.addFixup(ctx, fixup{
-		TreeKey:            treeKey,
+		// Clone the tree key, since we don't own the memory.
+		TreeKey:            slices.Clone(treeKey),
 		Type:               mergeFixup,
 		ParentPartitionKey: parentPartitionKey,
 		PartitionKey:       partitionKey,

pkg/sql/vecindex/cspann/index.go

@@ -350,6 +350,10 @@ func (vi *Index) Close() {
 // before Insert when a vector is updated. Even then, it's not guaranteed that
 // Delete will find the old vector. The search set handles this rare case by
 // filtering out results with duplicate key bytes.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
+// TODO(andyk): This is not true of the MemStore.
 func (vi *Index) Insert(
 	ctx context.Context, idxCtx *Context, treeKey TreeKey, vec vector.T, key KeyBytes,
 ) error {
@@ -393,6 +397,10 @@ func (vi *Index) Insert(
 //
 // NOTE: Even if the vector is removed, there may still be duplicate dangling
 // instances of the vector still remaining in the index.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
+// TODO(andyk): This is not true of the MemStore.
 func (vi *Index) Delete(
 	ctx context.Context, idxCtx *Context, treeKey TreeKey, vec vector.T, key KeyBytes,
 ) (deleted bool, err error) {
@@ -423,6 +431,10 @@ func (vi *Index) Delete(
 // and returns them in the search set. Set searchSet.MaxResults to limit the
 // number of results. This is called within the scope of a transaction so that
 // the index does not appear to change during the search.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
+// TODO(andyk): This is not true of the MemStore.
 func (vi *Index) Search(
 	ctx context.Context,
 	idxCtx *Context,
@@ -440,6 +452,10 @@ func (vi *Index) Search(
 // partition, as well as the centroid of the partition (in the Vector field).
 // This is useful for callers that directly insert KV rows rather than using
 // this library to do it.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
+// TODO(andyk): This is not true of the MemStore.
 func (vi *Index) SearchForInsert(
 	ctx context.Context, idxCtx *Context, treeKey TreeKey, vec vector.T,
 ) (*SearchResult, error) {
@@ -471,6 +487,10 @@ func (vi *Index) SearchForInsert(
 // It returns a single search result containing the key of that partition, or
 // nil if the vector cannot be found. This is useful for callers that directly
 // delete KV rows rather than using this library to do it.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
+// TODO(andyk): This is not true of the MemStore.
 func (vi *Index) SearchForDelete(
 	ctx context.Context, idxCtx *Context, treeKey TreeKey, vec vector.T, key KeyBytes,
 ) (*SearchResult, error) {

pkg/sql/vecindex/mutation_searcher.go

@@ -51,6 +51,9 @@ func (s *MutationSearcher) Init(idx *cspann.Index, txn *kv.Txn) {
 // SearchForInsert triggers a search for the partition in which to insert the
 // input vector. The partition's key is returned by PartitionKey() and the
 // input vector's quantized and encoded form is returned by EncodedVector().
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
 func (s *MutationSearcher) SearchForInsert(
 	ctx context.Context, prefix roachpb.Key, vec vector.T,
 ) error {
@@ -79,6 +82,9 @@ func (s *MutationSearcher) SearchForInsert(
 // SearchForDelete triggers a search for the partition which contains the vector
 // to be deleted, identified by its primary key. If the input vector is found,
 // its partition is returned by PartitionKey().
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
 func (s *MutationSearcher) SearchForDelete(
 	ctx context.Context, prefix roachpb.Key, vec vector.T, key cspann.KeyBytes,
 ) error {

pkg/sql/vecindex/searcher.go

@@ -61,6 +61,9 @@ func (s *Searcher) Init(idx *cspann.Index, txn *kv.Txn, baseBeamSize, maxResults
 // Search triggers a search over the index for the given vector, within the
 // scope of the given prefix. "maxResults" specifies the maximum number of
 // results that will be returned.
+//
+// NOTE: The caller is assumed to own the memory for all parameters and can
+// reuse the memory after the call returns.
 func (s *Searcher) Search(ctx context.Context, prefix roachpb.Key, vec vector.T) error {
 	err := s.idx.Search(ctx, &s.idxCtx, cspann.TreeKey(prefix), vec, &s.searchSet, s.options)
 	if err != nil {

pkg/sql/vecindex/searcher_test.go

@@ -7,6 +7,7 @@ package vecindex
 
 import (
 	"context"
+	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -89,22 +90,40 @@ func TestSearcher(t *testing.T) {
 	// Insert two vectors into root partition.
 	var mutator MutationSearcher
 	mutator.Init(idx, tx)
-	prefix := keys.MakeFamilyKey(encoding.EncodeVarintAscending([]byte{}, 100), 0 /* famID */)
+
+	// Reuse prefix, key bytes, value bytes and vector memory, to ensure it's
+	// allowed.
+	var prefix, keyBytes, valueBytes []byte
+	var original, randomized vector.T
 
 	insertVector := func(vec vector.T, key int64, val cspann.ValueBytes) vector.T {
+		prefix = prefix[:0]
+		keyBytes = keyBytes[:0]
+		randomized = randomized[:0]
+
+		prefix = encoding.EncodeVarintAscending(prefix, 100)
 		require.NoError(t, mutator.SearchForInsert(ctx, prefix, vec))
+
 		partitionKey := cspann.PartitionKey(*mutator.PartitionKey().(*tree.DInt))
-		keyBytes := keys.MakeFamilyKey(encoding.EncodeVarintAscending([]byte{}, key), 0 /* famID */)
-		randomizedVec := make(vector.T, len(vec))
-		idx.RandomizeVector(vec, randomizedVec)
+		keyBytes = keys.MakeFamilyKey(encoding.EncodeVarintAscending(keyBytes, key), 0 /* famID */)
+
+		randomized = slices.Grow(randomized, len(vec))[:len(vec)]
+		idx.RandomizeVector(vec, randomized)
 		err = mutator.txn.AddToPartition(ctx, cspann.TreeKey(prefix), partitionKey, cspann.LeafLevel,
-			randomizedVec, cspann.ChildKey{KeyBytes: keyBytes}, val)
+			randomized, cspann.ChildKey{KeyBytes: keyBytes}, val)
 		require.NoError(t, err)
-		return randomizedVec
+		return randomized
 	}
 
-	insertVector(vector.T{1, 2}, 1, cspann.ValueBytes{1, 2})
-	randomizedVec := insertVector(vector.T{5, 3}, 2, cspann.ValueBytes{3, 4})
+	// Reuse vector and value memory, to ensure it's allowed.
+	original = vector.T{1, 2}
+	valueBytes = []byte{1, 2}
+	insertVector(original, 1, cspann.ValueBytes(valueBytes))
+	original[0] = 5
+	original[1] = 3
+	valueBytes[0] = 3
+	valueBytes[1] = 4
+	randomized = insertVector(original, 2, cspann.ValueBytes(valueBytes))
 
 	// Validate that search vector was correctly encoded and quantized.
 	encodedVec := mutator.EncodedVector()
@@ -113,29 +132,52 @@ func TestSearcher(t *testing.T) {
 		[]byte(*encodedVec.(*tree.DBytes)), &vecSet)
 	require.NoError(t, err)
 	require.Empty(t, remainder)
-	require.Equal(t, randomizedVec, vecSet.Vectors.At(0))
+	require.Equal(t, randomized, vecSet.Vectors.At(0))
 
-	// Use the Searcher.
+	// Search for a vector that doesn't exist in the tree (reuse memory).
+	prefix = prefix[:0]
+	prefix = encoding.EncodeVarintAscending(prefix, 200)
+	original[0] = 1
+	original[1] = 1
 	var searcher Searcher
 	searcher.Init(idx, tx, 8 /* baseBeamSize */, 2 /* maxResults */)
-	require.NoError(t, searcher.Search(ctx, prefix, vector.T{1, 1}))
+	require.NoError(t, searcher.Search(ctx, prefix, original))
+	require.Nil(t, searcher.NextResult())
+
+	// Search for a vector that does exist (reuse memory).
+	prefix = prefix[:0]
+	prefix = encoding.EncodeVarintAscending(prefix, 100)
+	require.NoError(t, searcher.Search(ctx, prefix, original))
 	res := searcher.NextResult()
 	require.InDelta(t, float32(1), res.QuerySquaredDistance, 0.01)
 	res = searcher.NextResult()
 	require.InDelta(t, float32(20), res.QuerySquaredDistance, 0.01)
 	require.Nil(t, searcher.NextResult())
 
-	// Search for a vector to delete.
-	keyBytes := keys.MakeFamilyKey(encoding.EncodeVarintAscending([]byte{}, 1), 0 /* famID */)
-	require.NoError(t, mutator.SearchForDelete(ctx, prefix, vector.T{1, 2}, keyBytes))
-	require.Equal(t, tree.NewDInt(tree.DInt(1)), mutator.PartitionKey())
+	// Search for a vector to delete that doesn't exist (reuse memory).
+	keyBytes = keyBytes[:0]
+	original[0] = 1
+	original[1] = 2
+	keyBytes = keys.MakeFamilyKey(encoding.EncodeVarintAscending(keyBytes, 123), 0 /* famID */)
+	require.NoError(t, mutator.SearchForDelete(ctx, prefix, original, keyBytes))
+	require.Equal(t, tree.DNull, mutator.PartitionKey())
 	require.Nil(t, mutator.EncodedVector())
 
-	// Search for a vector to delete that doesn't exist.
-	keyBytes = keys.MakeFamilyKey(encoding.EncodeVarintAscending([]byte{}, 123), 0 /* famID */)
-	require.NoError(t, mutator.SearchForDelete(ctx, prefix, vector.T{1, 2}, keyBytes))
+	// Search for a vector to delete that doesn't exist in the tree (reuse memory).
+	prefix = prefix[:0]
+	prefix = encoding.EncodeVarintAscending(prefix, 200)
+	require.NoError(t, mutator.SearchForDelete(ctx, prefix, original, keyBytes))
 	require.Equal(t, tree.DNull, mutator.PartitionKey())
 	require.Nil(t, mutator.EncodedVector())
 
+	// Search for a vector to delete (reuse memory).
+	prefix = prefix[:0]
+	prefix = encoding.EncodeVarintAscending(prefix, 100)
+	keyBytes = keyBytes[:0]
+	keyBytes = keys.MakeFamilyKey(encoding.EncodeVarintAscending(keyBytes, 1), 0 /* famID */)
+	require.NoError(t, mutator.SearchForDelete(ctx, prefix, original, keyBytes))
+	require.Equal(t, tree.NewDInt(tree.DInt(1)), mutator.PartitionKey())
+	require.Nil(t, mutator.EncodedVector())
+
 	require.NoError(t, tx.Commit(ctx))
 }