Merge #148713

148713: cspann: move vectors from sibling partitions during split r=drewkimball a=andy-kimball

#### cspann: move vectors from sibling partitions during split

When a partition is split, the left and right target partitions are assigned
new centroids. This makes it possible for vectors in other partitions at the
same level to now be closer to one of those new centroids than they are to
their own centroid. In that case, we need to move those vectors to whichever
target partition is now closer.

#### cspann: enhance SearchSet

Add support in SearchSet for excluding partitions during search. Any
vectors in these partitions will not be added to the set. Also add
an option that includes the distance of each vector from its centroid
in search results. These options will be used by the split operation.

#### cspann: add BestCentroids vector index test

Add a new "best-centroids" test that prints out the partitions with the
closest centroids for a query vector. This is useful when gauging the
quality of the index.

Also, update the "recall" test to only sample from vectors that are not
part of the index. Achieving high recall is more challenging when
searching for such vectors.

#### cspann: move level searcher init code to Init method

Previously, the searcher initialized the search set for each level in
the Next method. This commit moves that code into the levelSearcher.Init
method. This makes some planned refactoring easier in an another commit.

#### cspann: add utils.ReplaceWithLast helper function

Add generic ReplaceWithLast function that removes an element from a
slice by replacing it with the last element and truncating the slice.

#### vecindex: add TryMoveVector method to Store interface

Add a new TryMoveVector to the Store interface. This method moves a
single vector from one partition to another as an atomic operation,
provided the destination metadata is as expected.

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

Author: craig[bot]

pkg/BUILD.bazel

@@ -658,6 +658,7 @@ ALL_TESTS = [
     "//pkg/sql/types:types_test",
     "//pkg/sql/vecindex/cspann/memstore:memstore_test",
     "//pkg/sql/vecindex/cspann/quantize:quantize_test",
+    "//pkg/sql/vecindex/cspann/utils:utils_test",
     "//pkg/sql/vecindex/cspann/workspace:workspace_test",
     "//pkg/sql/vecindex/cspann:cspann_test",
     "//pkg/sql/vecindex/vecencoding:vecencoding_test",
@@ -2384,6 +2385,7 @@ GO_TARGETS = [
     "//pkg/sql/vecindex/cspann/quantize:quantize_test",
     "//pkg/sql/vecindex/cspann/testutils:testutils",
     "//pkg/sql/vecindex/cspann/utils:utils",
+    "//pkg/sql/vecindex/cspann/utils:utils_test",
     "//pkg/sql/vecindex/cspann/workspace:workspace",
     "//pkg/sql/vecindex/cspann/workspace:workspace_test",
     "//pkg/sql/vecindex/cspann:cspann",

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -89,6 +89,7 @@ go_test(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_guptarohit_asciigraph//:asciigraph",
         "@com_github_stretchr_testify//require",
+        "@org_golang_x_exp//slices",
         "@org_gonum_v1_gonum//floats/scalar",
         "@org_gonum_v1_gonum//stat",
     ],

pkg/sql/vecindex/cspann/commontest/storetests.go

@@ -6,7 +6,9 @@
 package commontest
 
 import (
+	"cmp"
 	"context"
+	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -987,6 +989,135 @@ func (suite *StoreTestSuite) TestTryRemoveFromPartition() {
 	}
 }
 
+func (suite *StoreTestSuite) TestTryMoveVector() {
+	store := suite.makeStore(suite.quantizer)
+	defer store.Close(suite.T())
+
+	doTest := func(treeID int) {
+		treeKey := store.MakeTreeKey(suite.T(), treeID)
+
+		// Create source partition with some vectors.
+		sourcePartitionKey, _ := suite.createTestPartition(store, treeKey)
+
+		// Create empty target partition.
+		targetPartitionKey := cspann.PartitionKey(20)
+		metadata := cspann.PartitionMetadata{
+			Level:    cspann.SecondLevel,
+			Centroid: vector.T{2, 4},
+		}
+		metadata.StateDetails.MakeReady()
+		suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, targetPartitionKey, metadata))
+		targetPartition, err := store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+
+		// Source partition does not yet exist.
+		expected := *targetPartition.Metadata()
+		moved, err := store.TryMoveVector(
+			suite.ctx, treeKey, cspann.PartitionKey(99), targetPartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Destination partition does not yet exist.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, cspann.PartitionKey(99),
+			vec1, partitionKey1, valueBytes1, cspann.PartitionMetadata{})
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Source partition is the same as destination partition.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, sourcePartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Now move should work.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.True(moved)
+
+		// Fetch back the target partition and validate it.
+		targetPartition, err = store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey1}, targetPartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes1}, targetPartition.ValueBytes())
+
+		// Try to move again, but with mismatched expected metadata.
+		var errConditionFailed *cspann.ConditionFailedError
+		metadata = expected
+		metadata.StateDetails.State = cspann.DrainingForMergeState
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, metadata)
+		suite.ErrorAs(err, &errConditionFailed)
+		suite.False(moved)
+		suite.True(errConditionFailed.Actual.Equal(&expected))
+
+		// Try again, this time with correct expected metadata.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, expected)
+		suite.NoError(err)
+		suite.True(moved)
+
+		// Fetch back the source partition and validate it.
+		sourcePartition, err := store.TryGetPartition(suite.ctx, treeKey, sourcePartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey2}, sourcePartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes2}, sourcePartition.ValueBytes())
+
+		// Try to move a vector that no longer exists in the source partition.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Try to move a vector that already exists in the target partition.
+		added, err := store.TryAddToPartition(
+			suite.ctx, treeKey, targetPartitionKey, vec2.AsSet(),
+			[]cspann.ChildKey{partitionKey2}, []cspann.ValueBytes{valueBytes2}, expected)
+		suite.NoError(err)
+		suite.True(added)
+
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey2, valueBytes2, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Ensure that the vector was not removed from the source partition.
+		sourcePartition, err = store.TryGetPartition(suite.ctx, treeKey, sourcePartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey2}, sourcePartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes2}, sourcePartition.ValueBytes())
+
+		// Ensure that the target partition now has all three vectors.
+		targetPartition, err = store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+		suite.Equal(3, targetPartition.Count())
+		childKeys := slices.Clone(targetPartition.ChildKeys())
+		slices.SortFunc(childKeys, func(a, b cspann.ChildKey) int {
+			return cmp.Compare(a.PartitionKey, b.PartitionKey)
+		})
+		suite.Equal([]cspann.ChildKey{partitionKey1, partitionKey2, partitionKey3}, childKeys)
+	}
+
+	suite.Run("default tree", func() {
+		doTest(0)
+	})
+
+	if store.AllowMultipleTrees() {
+		// Ensure that vectors are independent across trees.
+		suite.Run("different tree", func() {
+			doTest(1)
+		})
+	}
+}
+
 func (suite *StoreTestSuite) TestTryClearPartition() {
 	store := suite.makeStore(suite.quantizer)
 	defer store.Close(suite.T())

pkg/sql/vecindex/cspann/fixup_processor.go

@@ -486,7 +486,8 @@ func (fp *FixupProcessor) nextFixup(ctx context.Context) (next fixup, ok bool) {
 				}()
 			}
 
-			if discard {
+			// Always process fixup if it's single-stepping.
+			if discard && !next.SingleStep {
 				fp.removeFixup(next)
 				continue
 			}