Merge #150573

150573: vecindex: fix and enhance vecindex testing r=drewkimball a=andy-kimball

#### vecindex: add datadriven tests for vecindex

Refactor the C-SPANN datadriven tests so they can be shared by other
Store implementations. Add new datadriven test for vecindex that
tests recall for several interesting datasets.

#### cspann: fix bug in CalculateTruth

Fix bug in CalculateTruth, which was calling vecpb.MeasureDistance with
un-normalized vectors for Cosine distance. This returns incorrect results.

#### vecindex: fix determinism violations

Fix a couple cases where the vector index is not behaving
deterministically even when configured to do so. In one case, partition
keys were not generated deterministically. In the other case, the seed
used in crdb_test builds was different than the seed used otherwise.

#### builtins: add process_vector_index_fixups builtin

Add a new builtin function that waits until all pending fixups have
been processed for a given index. This is used for testing.

Previously, we were processing fixups in deterministic tests by calling
ProcessFixups just before inserting vectors or searching for them.
However, that was causing deadlocks, because fixups were being processed
in the same transaction that was updating the index.

The vector_index test is now updated to call the new builtin when it
needs to force fixup processing.

#### cspann: use brackets to format vectors

Postgres uses brackets to format vectors, e.g. [1, 2]. However, our test
code currently uses parentheses, e.g. (1, 2). Update our tests to use
brackets instead, so that we're more consistent.

Co-authored-by: Andrew Kimball <[email protected]>

Author: craig[bot]

pkg/cmd/vecbench/mem_provider.go

@@ -177,6 +177,9 @@ func (m *MemProvider) Search(
 
 		// Get result keys.
 		results := searchSet.PopResults()
+		if len(results) > memState.maxResults {
+			results = results[:memState.maxResults]
+		}
 		keys = make([]cspann.KeyBytes, len(results))
 		for i, res := range results {
 			keys[i] = []byte(res.ChildKey.KeyBytes)
@@ -196,7 +199,9 @@ func (m *MemProvider) Save(ctx context.Context) error {
 	}
 
 	// Wait for any remaining background fixups to be processed.
-	m.index.ProcessFixups()
+	if err := m.index.ProcessFixups(ctx); err != nil {
+		return err
+	}
 
 	startTime := timeutil.Now()
 

pkg/sql/faketreeeval/evalctx.go

@@ -577,6 +577,13 @@ func (ep *DummyEvalPlanner) RetryCounter() int {
 	return 0
 }
 
+// ProcessVectorIndexFixups is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) ProcessVectorIndexFixups(
+	ctx context.Context, tableID descpb.ID, indexID descpb.IndexID,
+) error {
+	return nil
+}
+
 // DummyPrivilegedAccessor implements the tree.PrivilegedAccessor interface by returning errors.
 type DummyPrivilegedAccessor struct{}
 

pkg/sql/logictest/testdata/logic_test/vector_index

@@ -539,12 +539,37 @@ INSERT INTO exec_test (a, b, vec1) VALUES
   (2, 1, '[4, 5, 6]'),
   (3, 2, '[7, 8, 9]'),
   (4, 2, '[10, 11, 12]'),
-  (5, 2, '[13, 14, 15]'),
+  (5, 2, '[13, 14, 15]');
+
+statement ok
+INSERT INTO exec_test (a, b, vec1) VALUES
   (6, NULL, '[16, 17, 18]'),
   (7, NULL, '[1, 1, 1]'),
   (8, NULL, NULL),
   (9, 3, NULL);
 
+# Define helper to force fixup processing for the given index.
+statement ok
+CREATE PROCEDURE process_vector_index_fixups_for(desc_name STRING, idx_name STRING)
+LANGUAGE SQL
+AS $$
+  SELECT crdb_internal.process_vector_index_fixups(
+    (SELECT id FROM system.namespace WHERE name = desc_name),
+    (
+      SELECT index_id
+      FROM crdb_internal.table_indexes
+      WHERE descriptor_name = desc_name AND index_name = idx_name
+    )
+  )
+$$;
+
+# Process split fixups for the two indexes.
+statement ok
+CALL process_vector_index_fixups_for('exec_test', 'idx1');
+
+statement ok
+CALL process_vector_index_fixups_for('exec_test', 'idx2');
+
 statement error pgcode 22000 pq: expected 3 dimensions, not 1
 INSERT INTO exec_test (a, b, vec1) VALUES (10, 1, '[1]');
 
@@ -765,9 +790,22 @@ INSERT INTO distance_metrics (a, v) VALUES
   (1, '[0, 0]'),
   (2, '[-2, -2]'),
   (3, '[2, 2]'),
-  (4, '[4, 4]'),
+  (4, '[4, 4]');
+
+statement ok
+INSERT INTO distance_metrics (a, v) VALUES
   (5, '[-2, 2]');
 
+# Process split fixups for the three indexes.
+statement ok
+CALL process_vector_index_fixups_for('distance_metrics', 'idx1');
+
+statement ok
+CALL process_vector_index_fixups_for('distance_metrics', 'idx2');
+
+statement ok
+CALL process_vector_index_fixups_for('distance_metrics', 'idx3');
+
 # Results using L2 distance.
 query ITF rowsort
 SELECT a, v, round(v <-> '[2, 2]', 2) dist FROM distance_metrics ORDER BY v <-> '[2, 2]' LIMIT 5;

pkg/sql/planner.go

@@ -1077,3 +1077,14 @@ func (p *planner) ExtendHistoryRetention(ctx context.Context, jobID jobspb.JobID
 func (p *planner) RetryCounter() int {
 	return p.autoRetryCounter + p.autoRetryStmtCounter
 }
+
+// ProcessVectorIndexFixups is part of the eval.Planner interface.
+func (p *planner) ProcessVectorIndexFixups(
+	ctx context.Context, tableID descpb.ID, indexID descpb.IndexID,
+) error {
+	vi, err := p.execCfg.VecIndexManager.Get(ctx, tableID, indexID)
+	if err != nil {
+		return err
+	}
+	return vi.ProcessFixups(ctx)
+}

pkg/sql/sem/builtins/builtins.go

@@ -9357,6 +9357,30 @@ WHERE object_id = table_descriptor_id
 			},
 		},
 	),
+	"crdb_internal.process_vector_index_fixups": makeBuiltin(
+		tree.FunctionProperties{
+			Category:     builtinconstants.CategoryTesting,
+			Undocumented: true,
+		},
+		tree.Overload{
+			Types: tree.ParamTypes{
+				{Name: "table_id", Typ: types.Int},
+				{Name: "index_id", Typ: types.Int},
+			},
+			ReturnType: tree.FixedReturnType(types.Void),
+			Fn: func(ctx context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				tableID := descpb.ID(tree.MustBeDInt(args[0]))
+				indexID := descpb.IndexID(tree.MustBeDInt(args[1]))
+				err := evalCtx.Planner.ProcessVectorIndexFixups(ctx, tableID, indexID)
+				if err != nil {
+					return nil, err
+				}
+				return tree.DVoidDatum, nil
+			},
+			Info:       "Waits until all outstanding fixups for the vector index with the given ID have been processed.",
+			Volatility: volatility.Volatile,
+		},
+	),
 }
 
 var lengthImpls = func(incBitOverload bool) builtinDefinition {

pkg/sql/sem/builtins/fixed_oids.go

@@ -2795,6 +2795,7 @@ var builtinOidsArray = []string{
 	2840: `lead(val: ltree, n: int) -> ltree`,
 	2841: `lead(val: ltree, n: int, default: ltree) -> ltree`,
 	2842: `last_value(val: ltree) -> ltree`,
+	2843: `crdb_internal.process_vector_index_fixups(table_id: int, index_id: int) -> void`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/eval/deps.go

@@ -455,6 +455,10 @@ type Planner interface {
 
 	// RetryCounter is the number of times this statement has been retried.
 	RetryCounter() int
+
+	// ProcessVectorIndexFixups waits until all outstanding fixups for the vector
+	// index with the given ID have been processed.
+	ProcessVectorIndexFixups(ctx context.Context, tableID descpb.ID, indexID descpb.IndexID) error
 }
 
 // InternalRows is an iterator interface that's exposed by the internal

pkg/sql/vecindex/BUILD.bazel

@@ -32,7 +32,6 @@ go_library(
         "//pkg/sql/vecindex/vecpb",
         "//pkg/sql/vecindex/vecstore",
         "//pkg/sql/vecindex/vecstore/vecstorepb",
-        "//pkg/util/buildutil",
         "//pkg/util/errorutil/unimplemented",
         "//pkg/util/log",
         "//pkg/util/metric",
@@ -51,7 +50,7 @@ go_test(
         "searcher_test.go",
         "vecindex_test.go",
     ],
-    data = ["//pkg/sql/vecindex/cspann:datasets"],
+    data = ["//pkg/sql/vecindex/cspann:datasets"] + glob(["testdata/**"]),
     embed = [":vecindex"],
     exec_properties = select({
         "//build/toolchains:is_heavy": {"test.Pool": "large"},
@@ -83,6 +82,7 @@ go_test(
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
         "//pkg/sql/vecindex/cspann",
+        "//pkg/sql/vecindex/cspann/commontest",
         "//pkg/sql/vecindex/cspann/quantize",
         "//pkg/sql/vecindex/cspann/testutils",
         "//pkg/sql/vecindex/vecencoding",
@@ -98,6 +98,7 @@ go_test(
         "//pkg/util/log",
         "//pkg/util/randutil",
         "//pkg/util/vector",
+        "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/vecindex/cspann/commontest/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 go_library(
     name = "commontest",
     srcs = [
+        "indextests.go",
         "storetests.go",
         "utils.go",
     ],
@@ -16,6 +17,7 @@ go_library(
         "//pkg/sql/vecindex/vecpb",
         "//pkg/util/encoding",
         "//pkg/util/vector",
+        "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
         "@com_github_stretchr_testify//suite",