backfill: handle transaction retry for vector index backfill

Previously, when backfilling a vector index, a transaction retry would
cause backfill failure because the vector index backfill writer would
overwrite the values read by the backfill reader with what it wanted to
push down to KV so as to avoid new allocations. This worked well so long
as there were no transaction retries but would fail to re-encode the
index entry on a retry because the writer lost access to the unquantized
vector. The quantized vector cannot be reused on a transaction retry
because fixups may cause the target partition to change.

This patch creates a scratch rowenc.IndexEntry in the backfill helper to
store the input vector entry. Before attempting to write the entry, we
copy the input IndexEntry to the scratch entry and use that to re-encode
the vector, which is still modified in place to limit new allocations.

Additionally, this patch switches the writer from using CPut() to
CPutAllowingIfNotExists() so that if the backfill job restarts, we don't
see the partially written index and fail due to duplicate keys.

Informs: #143107
Release note: None

Author: mw5h

Diff for: pkg/sql/backfill/BUILD.bazel

@@ -60,19 +60,30 @@ go_library(
 go_test(
     name = "backfill_test",
     srcs = [
+        "backfill_test.go",
         "index_backfiller_cols_test.go",
         "mvcc_index_merger_test.go",
     ],
     embed = [":backfill"],
     deps = [
+        "//pkg/base",
         "//pkg/keys",
         "//pkg/kv",
         "//pkg/roachpb",
+        "//pkg/security/securityassets",
+        "//pkg/security/securitytest",
+        "//pkg/server",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/catenumpb",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/tabledesc",
+        "//pkg/sql/execinfra",
         "//pkg/sql/sem/catid",
+        "//pkg/testutils/serverutils",
+        "//pkg/testutils/sqlutils",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/randutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],

Diff for: pkg/sql/backfill/backfill.go

@@ -485,17 +485,19 @@ type VectorIndexHelper struct {
 // context of the provided transaction. This lookup then gives the leaf partition
 // where the index entry is to be inserted and the unquantized vector can then be
 // re-encoded to get the properly quantized vector with the new partition's
-// centroid.
+// centroid. The new key is then returned in the outputEntry. The inputEntry is
+// not overwritten in case the transaction has to be retried.
 func (vih *VectorIndexHelper) ReEncodeVector(
-	ctx context.Context, txn *kv.Txn, keyBytes []byte, entry *rowenc.IndexEntry,
+	ctx context.Context, txn *kv.Txn, inputEntry rowenc.IndexEntry, outputEntry *rowenc.IndexEntry,
 ) (*rowenc.IndexEntry, error) {
+	keyBytes := inputEntry.Key[len(vih.indexPrefix):]
 	key, err := vecencoding.DecodeVectorKey(keyBytes, vih.numPrefixCols)
 	if err != nil {
 		return &rowenc.IndexEntry{}, err
 	}
 
 	// Decode vector and suffix bytes from the entry value.
-	valueBytes, err := entry.Value.GetBytes()
+	valueBytes, err := inputEntry.Value.GetBytes()
 	if err != nil {
 		return &rowenc.IndexEntry{}, err
 	}
@@ -517,16 +519,16 @@ func (vih *VectorIndexHelper) ReEncodeVector(
 		panic("expected encoded vector to be of type DBytes")
 	}
 
-	entry.Key = entry.Key[:0]
-	entry.Key = append(entry.Key, vih.indexPrefix...)
-	entry.Key = key.Encode(entry.Key)
+	outputEntry.Key = append(outputEntry.Key[:0], vih.indexPrefix...)
+	outputEntry.Key = key.Encode(outputEntry.Key)
+
 	entryValueLen := vecencoding.EncodedVectorValueLen(quantizedVector.UnsafeBytes(), suffix)
-	buf := entry.Value.AllocBytes(entryValueLen)[:0]
+	buf := outputEntry.Value.AllocBytes(entryValueLen)[:0]
 	vecencoding.EncodeVectorValue(buf, quantizedVector.UnsafeBytes(), suffix)
 
-	// entry.Family is left unchanged from the entry we received, originally encoded by rowenc.EncodeSecondaryIndexes()
+	outputEntry.Family = inputEntry.Family
 
-	return entry, nil
+	return outputEntry, nil
 }
 
 // IndexBackfiller is capable of backfilling all the added index.

Diff for: pkg/sql/backfill/backfill_test.go

@@ -0,0 +1,106 @@
+package backfill_test
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/security/securityassets"
+	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMain(m *testing.M) {
+	securityassets.SetLoader(securitytest.EmbeddedAssets)
+	randutil.SeedForTests()
+	serverutils.InitTestServerFactory(server.TestServerFactory)
+
+	os.Exit(m.Run())
+}
+
+// testRetryError implements pgerror.ClientVisibleRetryError
+type testRetryError struct{}
+
+func (e *testRetryError) Error() string {
+	return "test retry error"
+}
+
+func (e *testRetryError) ClientVisibleRetryError() {}
+
+func TestVectorColumnAndIndexBackfill(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// Track whether we've injected an error
+	var errorState struct {
+		mu         syncutil.Mutex
+		hasErrored bool
+	}
+
+	ctx := context.Background()
+	srv, db, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			DistSQL: &execinfra.TestingKnobs{
+				// Inject a retriable error on the first call to the vector index backfiller.
+				VectorIndexBackfillTxnError: func() error {
+					errorState.mu.Lock()
+					defer errorState.mu.Unlock()
+					if !errorState.hasErrored {
+						errorState.hasErrored = true
+						return &testRetryError{}
+					}
+					return nil
+				},
+			},
+		},
+	})
+	defer srv.Stopper().Stop(ctx)
+	sqlDB := sqlutils.MakeSQLRunner(db)
+
+	// Create a table with a vector column
+	sqlDB.Exec(t, `
+		CREATE TABLE vectors (
+			id INT PRIMARY KEY,
+			vec VECTOR(3)
+		)
+	`)
+
+	// Insert 200 rows with random vector data
+	sqlDB.Exec(t, `
+		INSERT INTO vectors (id, vec)
+		SELECT 
+			generate_series(1, 200) as id,
+			ARRAY[random(), random(), random()]::vector(3) as vec
+	`)
+
+	// Create a vector index on the vector column
+	sqlDB.Exec(t, `
+		CREATE VECTOR INDEX vec_idx ON vectors (vec)
+	`)
+
+	// Test vector similarity search and see that the backfiller got at
+	// least some of the vectors in there.
+	var matchCount int
+	sqlDB.QueryRow(t, `
+		SELECT count(*)
+		FROM (
+			SELECT id 
+			FROM vectors@vec_idx
+			ORDER BY vec <-> ARRAY[0.5, 0.5, 0.5]::vector(3)
+			LIMIT 200 
+		)
+	`).Scan(&matchCount)
+	// There's some non-determinism here where we may not find all 200 vectors.
+	// I chose 190 as a low water mark to prevent test flakes, but it should really
+	// be 200 in most cases.
+	require.Greater(t, matchCount, 190, "Expected to find at least 190 similar vectors")
+}

Diff for: pkg/sql/execinfra/server_config.go

@@ -259,6 +259,10 @@ type TestingKnobs struct {
 	// indexes.
 	IndexBackfillProgressReportInterval time.Duration
 
+	// VectorIndexBackfillTxnError is called during vector index entry backfill to
+	// simulate a transaction error.
+	VectorIndexBackfillTxnError func() error
+
 	// ForceDiskSpill forces any processors/operators that can fall back to disk
 	// to fall back to disk immediately.
 	//

Diff for: pkg/sql/rowexec/indexbackfiller.go

@@ -182,9 +182,9 @@ func (ib *indexBackfiller) constructIndexEntries(
 }
 
 func (ib *indexBackfiller) maybeReencodeVectorIndexEntry(
-	ctx context.Context, indexEntry *rowenc.IndexEntry,
+	ctx context.Context, tmpEntry *rowenc.IndexEntry, indexEntry *rowenc.IndexEntry,
 ) (bool, error) {
-	indexID, keyBytes, err := rowenc.DecodeIndexKeyPrefix(ib.flowCtx.EvalCtx.Codec, ib.desc.GetID(), indexEntry.Key)
+	indexID, _, err := rowenc.DecodeIndexKeyPrefix(ib.flowCtx.EvalCtx.Codec, ib.desc.GetID(), indexEntry.Key)
 	if err != nil {
 		return false, err
 	}
@@ -194,14 +194,35 @@ func (ib *indexBackfiller) maybeReencodeVectorIndexEntry(
 		return false, nil
 	}
 
+	// Initialize the tmpEntry. This will store the input entry that we are encoding
+	// so that we can overwrite the indexEntry we were given with the output
+	// encoding. This allows us to preserve the initial template indexEntry across
+	// transaction retries.
+	if tmpEntry.Key == nil {
+		tmpEntry.Key = make(roachpb.Key, len(indexEntry.Key))
+		tmpEntry.Value.RawBytes = make([]byte, len(indexEntry.Value.RawBytes))
+	}
+	tmpEntry.Key = append(tmpEntry.Key[:0], indexEntry.Key...)
+	tmpEntry.Value.RawBytes = append(tmpEntry.Value.RawBytes[:0], indexEntry.Value.RawBytes...)
+	tmpEntry.Family = indexEntry.Family
+
 	err = ib.flowCtx.Cfg.DB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
-		entry, err := vih.ReEncodeVector(ctx, txn.KV(), keyBytes, indexEntry)
+		entry, err := vih.ReEncodeVector(ctx, txn.KV(), *tmpEntry, indexEntry)
 		if err != nil {
 			return err
 		}
 
+		if ib.flowCtx.Cfg.TestingKnobs.VectorIndexBackfillTxnError != nil {
+			err = ib.flowCtx.Cfg.TestingKnobs.VectorIndexBackfillTxnError()
+			if err != nil {
+				return err
+			}
+		}
+
 		b := txn.KV().NewBatch()
-		b.CPut(entry.Key, &entry.Value, nil)
+		// If the backfill job was interrupted and restarted, we may redo work, so allow
+		// that we may see duplicate keys here.
+		b.CPutAllowingIfNotExists(entry.Key, &entry.Value, entry.Value.TagAndDataBytes())
 		return txn.KV().Run(ctx, b)
 	})
 	if err != nil {
@@ -289,6 +310,7 @@ func (ib *indexBackfiller) ingestIndexEntries(
 	g.GoCtx(func(ctx context.Context) error {
 		defer close(stopProgress)
 
+		var vectorInputEntry rowenc.IndexEntry
 		for indexBatch := range indexEntryCh {
 			for _, indexEntry := range indexBatch.indexEntries {
 				// If there is at least one vector index being written, we need to check to see
@@ -298,7 +320,7 @@ func (ib *indexBackfiller) ingestIndexEntries(
 				// TODO(mw5h): batch up multiple index entries into a single batch.
 				// As is, we insert a single vector per batch, which is very slow.
 				if len(ib.VectorIndexes) > 0 {
-					isVectorIndex, err := ib.maybeReencodeVectorIndexEntry(ctx, &indexEntry)
+					isVectorIndex, err := ib.maybeReencodeVectorIndexEntry(ctx, &vectorInputEntry, &indexEntry)
 					if err != nil {
 						return err
 					} else if isVectorIndex {