Merge #142793

142793: vecindex: check partition size in fixup processor r=andy-kimball a=andy-kimball

Previously, we checked the size of a partition after adding or removing a vector to/from it. The main reason for doing this was to decide whether the partition needed to be split or merged. However, this added to the contention footprint of the user's transaction. This commit moves the partition size check to the background processor, as part of the SplitOrMergeFixup, and makes the size scan INCONSISTENT. While this can return stale data, it's OK if a split or merge gets delayed because of that. If the inconsistent check indicates that a split or merge may be necessary, we run a consistent scan to ensure that it's really necessary.

Besides being used to check for splits/merges, the partition size scan was used to ensure that we don't create an unbalanced K-means tree by moving the last vector of a non-root partition to a sibling partition. Now that RemoveFromPartition no longer returns a count, we instead use a new "skipLonelyVector" option when searching. This option stops the search from returning any vectors that are last in their partition, which prevents the creation of empty non-root partitions.

Epic: CRDB-42943

Release note: None

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

Author: craig[bot]

pkg/sql/vecindex/BUILD.bazel

@@ -39,8 +39,9 @@ go_test(
     srcs = [
         "manager_test.go",
         "searcher_test.go",
+        "vecindex_test.go",
     ],
-    data = glob(["testdata/**"]),
+    data = ["//pkg/sql/vecindex/cspann:features_10000"],
     embed = [":vecindex"],
     deps = [
         "//pkg/base",
@@ -63,6 +64,7 @@ go_test(
         "//pkg/sql/types",
         "//pkg/sql/vecindex/cspann",
         "//pkg/sql/vecindex/cspann/quantize",
+        "//pkg/sql/vecindex/cspann/testutils",
         "//pkg/sql/vecindex/vecpb",
         "//pkg/sql/vecindex/vecstore",
         "//pkg/testutils/serverutils",

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -3,9 +3,9 @@ load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 filegroup(
-    name = "testdata",
-    srcs = glob(["testdata/**"]),
-    visibility = ["//pkg/sql/vecindex/cspann:__subpackages__"],
+    name = "features_10000",
+    srcs = glob(["testdata/features_10000.gob"]),
+    visibility = ["//visibility:public"],
 )
 
 go_library(
@@ -55,7 +55,7 @@ go_test(
         "partition_test.go",
         "search_set_test.go",
     ],
-    data = ["//pkg/sql/vecindex/cspann:testdata"],
+    data = glob(["testdata/**"]),
     embed = [":cspann"],
     deps = [
         "//pkg/sql/vecindex/cspann/commontest",

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

@@ -18,11 +18,24 @@ import (
 
 // CheckPartitionMetadata tests the correctness of the given metadata's fields.
 func CheckPartitionMetadata(
-	t *testing.T, metadata cspann.PartitionMetadata, level cspann.Level, centroid vector.T, count int,
+	t *testing.T, metadata cspann.PartitionMetadata, level cspann.Level, centroid vector.T,
 ) {
 	require.Equal(t, level, metadata.Level)
 	require.Equal(t, []float32(centroid), testutils.RoundFloats(metadata.Centroid, 2))
-	require.Equal(t, count, metadata.Count)
+}
+
+// CheckPartitionCount tests the size of a partition.
+func CheckPartitionCount(
+	ctx context.Context,
+	t *testing.T,
+	store cspann.Store,
+	treeKey cspann.TreeKey,
+	partitionKey cspann.PartitionKey,
+	expectedCount int,
+) {
+	count, err := store.EstimatePartitionCount(ctx, treeKey, partitionKey)
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, count)
 }
 
 // BeginTransaction starts a new transaction for the given store and returns it.

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

@@ -23,9 +23,9 @@ import (
 type fixupType int
 
 const (
-	// splitOrMergeFixup is a fixup that includes the key of a partition to
-	// split or merge as well as the key of its parent partition (if it exists).
-	// Whether a partition is split or merged depends on its size.
+	// splitOrMergeFixup is a fixup that includes the key of a partition that
+	// may need to be split or merged, as well as the key of its parent partition
+	// (if it exists). Whether a partition is split or merged depends on its size.
 	splitOrMergeFixup fixupType = iota + 1
 	// vectorDeleteFixup is a fixup that includes the primary key of a vector to
 	// delete from the index, as well as the key of the partition that contains
@@ -266,20 +266,9 @@ func (fp *FixupProcessor) DelayInsertOrDelete(ctx context.Context) error {
 	return nil
 }
 
-// AddSplit enqueues a split fixup for later processing.
-func (fp *FixupProcessor) AddSplit(
-	ctx context.Context, treeKey TreeKey, parentPartitionKey PartitionKey, partitionKey PartitionKey,
-) {
-	fp.addFixup(ctx, fixup{
-		TreeKey:            treeKey,
-		Type:               splitOrMergeFixup,
-		ParentPartitionKey: parentPartitionKey,
-		PartitionKey:       partitionKey,
-	})
-}
-
-// AddMerge enqueues a merge fixup for later processing.
-func (fp *FixupProcessor) AddMerge(
+// AddSplitOrMergeCheck enqueues a fixup to check whether a split or merge is
+// needed for the given partition.
+func (fp *FixupProcessor) AddSplitOrMergeCheck(
 	ctx context.Context, treeKey TreeKey, parentPartitionKey PartitionKey, partitionKey PartitionKey,
 ) {
 	fp.addFixup(ctx, fixup{

pkg/sql/vecindex/cspann/fixup_processor.go

@@ -105,6 +105,18 @@ func (fw *fixupWorker) Start(ctx context.Context) {
 func (fw *fixupWorker) splitOrMergePartition(
 	ctx context.Context, parentPartitionKey PartitionKey, partitionKey PartitionKey,
 ) (err error) {
+	// Do a quick, inconsistent scan of the partition, in order to see if it may
+	// need to be split or merged.
+	count, err := fw.index.store.EstimatePartitionCount(ctx, fw.treeKey, partitionKey)
+	if err != nil {
+		return errors.Wrapf(err, "counting vectors in partition %d", partitionKey)
+	}
+	split := count > fw.index.options.MaxPartitionSize
+	merge := partitionKey != RootKey && count < fw.index.options.MinPartitionSize
+	if !split && !merge {
+		return nil
+	}
+
 	// Run the split or merge within a transaction.
 	fw.txn, err = fw.index.store.BeginTransaction(ctx)
 	if err != nil {
@@ -128,9 +140,10 @@ func (fw *fixupWorker) splitOrMergePartition(
 		return errors.Wrapf(err, "getting partition %d to split or merge", partitionKey)
 	}
 
-	// Don't split or merge the partition if its size is within bounds.
-	split := partition.Count() > fw.index.options.MaxPartitionSize
-	merge := partition.Count() < fw.index.options.MinPartitionSize
+	// Re-check the size of the partition now that it's locked, using a consistent
+	// scan, so that we are not acting based on stale information.
+	split = partition.Count() > fw.index.options.MaxPartitionSize
+	merge = partitionKey != RootKey && partition.Count() < fw.index.options.MinPartitionSize
 	if !split && !merge {
 		log.VEventf(ctx, 2, "partition %d size is within bounds, do not split or merge", partitionKey)
 		return nil
@@ -213,14 +226,14 @@ func (fw *fixupWorker) splitPartition(
 				partitionKey, parentPartitionKey)
 		}
 
-		metadata, err := fw.index.removeFromPartition(
-			ctx, fw.txn, fw.treeKey, parentPartitionKey, childKey)
+		err := fw.index.removeFromPartition(ctx, fw.txn, fw.treeKey, parentPartitionKey, childKey)
 		if err != nil {
 			return errors.Wrapf(err, "removing splitting partition %d from its parent %d",
 				partitionKey, parentPartitionKey)
 		}
 
-		if metadata.Count != 0 {
+		// Only attempt to move vectors if there are siblings.
+		if parentPartition.Count() > 1 {
 			// Move any vectors to sibling partitions that have closer centroids.
 			// Lazily get parent vectors only if they're actually needed.
 			var parentVectors vector.Set
@@ -474,8 +487,8 @@ func (fw *fixupWorker) moveVectorsToSiblings(
 		// there instead.
 		childKey := split.Partition.ChildKeys()[i]
 		valueBytes := split.Partition.ValueBytes()[i]
-		err = fw.index.addToPartition(ctx, fw.txn, fw.treeKey,
-			parentPartitionKey, siblingPartitionKey, vector, childKey, valueBytes)
+		err = fw.index.addToPartition(
+			ctx, fw.txn, fw.treeKey, siblingPartitionKey, vector, childKey, valueBytes)
 		if err != nil {
 			return errors.Wrapf(err, "moving vector to partition %d", siblingPartitionKey)
 		}
@@ -504,6 +517,11 @@ func (fw *fixupWorker) linkNearbyVectors(
 	idxCtx.level = partition.Level()
 	idxCtx.randomized = partition.Centroid()
 
+	// Ensure that the search never returns the last remaining vector in a
+	// non-leaf partition, in order to avoid moving it and creating an empty
+	// non-leaf partition, which is not allowed by a balanced K-means tree.
+	idxCtx.ignoreLonelyVector = partition.Level() != LeafLevel
+
 	// Don't link more vectors than the number of remaining slots in the split
 	// partition, to avoid triggering another split.
 	maxResults := fw.index.options.MaxPartitionSize - partition.Count()
@@ -542,24 +560,12 @@ func (fw *fixupWorker) linkNearbyVectors(
 		}
 
 		// Remove the vector from the other partition.
-		metadata, err := fw.index.removeFromPartition(
+		err = fw.index.removeFromPartition(
 			ctx, fw.txn, fw.treeKey, result.ParentPartitionKey, result.ChildKey)
 		if err != nil {
 			return errors.Wrapf(err, "removing vector from nearby partition %d during split of %d",
 				result.ParentPartitionKey, oldPartitionKey)
 		}
-		if metadata.Count == 0 && partition.Level() > LeafLevel {
-			// Removing the vector will result in an empty non-leaf partition, which
-			// is not allowed, as the K-means tree would not be fully balanced. Add
-			// the vector back to the partition. This is a very rare case and that
-			// partition is likely to be merged away regardless.
-			_, err = fw.txn.AddToPartition(
-				ctx, fw.treeKey, result.ParentPartitionKey, vector, result.ChildKey, result.ValueBytes)
-			if err != nil {
-				return errors.Wrapf(err, "adding vector to splitting partition %d", oldPartitionKey)
-			}
-			continue
-		}
 
 		// Add the vector to the split partition.
 		partition.Add(&fw.workspace, vector, result.ChildKey, result.ValueBytes)
@@ -627,7 +633,7 @@ func (fw *fixupWorker) mergePartition(
 			partitionKey, parentPartitionKey)
 		return nil
 	}
-	_, err = fw.index.removeFromPartition(ctx, fw.txn, fw.treeKey, parentPartitionKey, childKey)
+	err = fw.index.removeFromPartition(ctx, fw.txn, fw.treeKey, parentPartitionKey, childKey)
 	if err != nil {
 		return errors.Wrapf(err, "remove partition %d from parent partition %d",
 			partitionKey, parentPartitionKey)
@@ -687,7 +693,7 @@ func (fw *fixupWorker) deleteVector(
 		return nil
 	}
 
-	_, err = fw.index.removeFromPartition(ctx, fw.txn, fw.treeKey, partitionKey, childKey)
+	err = fw.index.removeFromPartition(ctx, fw.txn, fw.treeKey, partitionKey, childKey)
 	if errors.Is(err, ErrPartitionNotFound) {
 		log.VEventf(ctx, 2, "partition %d no longer exists, do not delete vector", partitionKey)
 		return nil

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -115,6 +115,11 @@ type Context struct {
 	// randomized is the original vector after it has been randomized by applying
 	// a random orthogonal transformation (ROT).
 	randomized vector.T
+	// ignoreLonelyVector, if true, prohibits searches from returning a vector
+	// that is the last remaining in its partition. This is used to avoid moving
+	// the last remaining vector to another partition, thereby creating an empty
+	// non-leaf partition, which is not allowed in a balanced K-means tree.
+	ignoreLonelyVector bool
 
 	tempSearchSet       SearchSet
 	tempSubSearchSet    SearchSet
@@ -364,9 +369,8 @@ func (vi *Index) Delete(
 	}
 
 	// Remove the vector from its partition in the store.
-	_, err = vi.removeFromPartition(
+	return vi.removeFromPartition(
 		ctx, idxCtx.txn, treeKey, result.ParentPartitionKey, result.ChildKey)
-	return err
 }
 
 // Search finds vectors in the index that are closest to the given query vector
@@ -404,16 +408,12 @@ func (vi *Index) SearchForInsert(
 		return nil, err
 	}
 
-	// Now fetch the centroid of the insert partition. This has the side effect
-	// of checking the size of the partition, in case it's over-sized.
+	// Now fetch the centroid of the insert partition.
 	partitionKey := result.ChildKey.PartitionKey
 	metadata, err := idxCtx.txn.GetPartitionMetadata(ctx, treeKey, partitionKey, true /* forUpdate */)
 	if err != nil {
 		return nil, err
 	}
-	if metadata.Count > vi.options.MaxPartitionSize {
-		vi.fixups.AddSplit(ctx, treeKey, result.ParentPartitionKey, partitionKey)
-	}
 
 	result.Vector = metadata.Centroid
 	return result, nil
@@ -474,18 +474,11 @@ func (vi *Index) ProcessFixups() {
 	vi.fixups.Process()
 }
 
-// ForceSplit enqueues a split fixup. It is used for testing.
-func (vi *Index) ForceSplit(
+// ForceSplitOrMerge enqueues a split or merge fixup. It is used for testing.
+func (vi *Index) ForceSplitOrMerge(
 	ctx context.Context, treeKey TreeKey, parentPartitionKey PartitionKey, partitionKey PartitionKey,
 ) {
-	vi.fixups.AddSplit(ctx, treeKey, parentPartitionKey, partitionKey)
-}
-
-// ForceMerge enqueues a merge fixup. It is used for testing.
-func (vi *Index) ForceMerge(
-	ctx context.Context, treeKey TreeKey, parentPartitionKey PartitionKey, partitionKey PartitionKey,
-) {
-	vi.fixups.AddMerge(ctx, treeKey, parentPartitionKey, partitionKey)
+	vi.fixups.AddSplitOrMergeCheck(ctx, treeKey, parentPartitionKey, partitionKey)
 }
 
 // setupInsertContext sets up the given context for an insert operation. Before
@@ -526,10 +519,9 @@ func (vi *Index) insertHelper(
 	if err != nil {
 		return err
 	}
-	parentPartitionKey := result.ParentPartitionKey
 	partitionKey := result.ChildKey.PartitionKey
-	err = vi.addToPartition(ctx, idxCtx.txn, idxCtx.treeKey, parentPartitionKey,
-		partitionKey, idxCtx.randomized, childKey, valueBytes)
+	err = vi.addToPartition(
+		ctx, idxCtx.txn, idxCtx.treeKey, partitionKey, idxCtx.randomized, childKey, valueBytes)
 	if errors.Is(err, ErrRestartOperation) {
 		return vi.insertHelper(ctx, idxCtx, childKey, valueBytes)
 	}
@@ -551,7 +543,11 @@ func (vi *Index) searchForInsertHelper(
 		return nil, errors.AssertionFailedf(
 			"SearchForInsert should return exactly one result, got %d", len(results))
 	}
-	return &results[0], err
+
+	vi.fixups.AddSplitOrMergeCheck(
+		ctx, idxCtx.treeKey, results[0].ParentPartitionKey, results[0].ChildKey.PartitionKey)
+
+	return &results[0], nil
 }
 
 // addToPartition calls the store to add the given vector to an existing
@@ -561,36 +557,31 @@ func (vi *Index) addToPartition(
 	ctx context.Context,
 	txn Txn,
 	treeKey TreeKey,
-	parentPartitionKey PartitionKey,
 	partitionKey PartitionKey,
 	vec vector.T,
 	childKey ChildKey,
 	valueBytes ValueBytes,
 ) error {
-	metadata, err := txn.AddToPartition(ctx, treeKey, partitionKey, vec, childKey, valueBytes)
+	err := txn.AddToPartition(ctx, treeKey, partitionKey, vec, childKey, valueBytes)
 	if err != nil {
 		return errors.Wrapf(err, "adding vector to partition %d", partitionKey)
 	}
-	if metadata.Count > vi.options.MaxPartitionSize {
-		vi.fixups.AddSplit(ctx, treeKey, parentPartitionKey, partitionKey)
-	}
 	return vi.stats.OnAddOrRemoveVector(ctx)
 }
 
 // removeFromPartition calls the store to remove a vector, by its key, from an
 // existing partition.
 func (vi *Index) removeFromPartition(
 	ctx context.Context, txn Txn, treeKey TreeKey, partitionKey PartitionKey, childKey ChildKey,
-) (metadata PartitionMetadata, err error) {
-	metadata, err = txn.RemoveFromPartition(ctx, treeKey, partitionKey, childKey)
+) error {
+	err := txn.RemoveFromPartition(ctx, treeKey, partitionKey, childKey)
 	if err != nil {
-		return PartitionMetadata{},
-			errors.Wrapf(err, "removing vector from partition %d", partitionKey)
+		return errors.Wrapf(err, "removing vector from partition %d", partitionKey)
 	}
-	if err := vi.stats.OnAddOrRemoveVector(ctx); err != nil {
-		return PartitionMetadata{}, err
+	if err = vi.stats.OnAddOrRemoveVector(ctx); err != nil {
+		return err
 	}
-	return metadata, nil
+	return nil
 }
 
 // searchHelper contains the core search logic for the K-means tree. It begins
@@ -768,11 +759,19 @@ func (vi *Index) searchChildPartitions(
 		count := idxCtx.tempCounts[i]
 		searchSet.Stats.SearchedPartition(level, count)
 
+		// If one of the searched partitions has only 1 vector remaining, do not
+		// return that vector when "ignoreLonelyVector" is true.
+		if idxCtx.ignoreLonelyVector && idxCtx.level == level && count == 1 {
+			searchSet.RemoveResults(parentResults[i].ChildKey.PartitionKey)
+		}
+
 		partitionKey := parentResults[i].ChildKey.PartitionKey
 		if count < vi.options.MinPartitionSize && partitionKey != RootKey {
-			vi.fixups.AddMerge(ctx, idxCtx.treeKey, parentResults[i].ParentPartitionKey, partitionKey)
+			vi.fixups.AddSplitOrMergeCheck(
+				ctx, idxCtx.treeKey, parentResults[i].ParentPartitionKey, partitionKey)
 		} else if count > vi.options.MaxPartitionSize {
-			vi.fixups.AddSplit(ctx, idxCtx.treeKey, parentResults[i].ParentPartitionKey, partitionKey)
+			vi.fixups.AddSplitOrMergeCheck(
+				ctx, idxCtx.treeKey, parentResults[i].ParentPartitionKey, partitionKey)
 		}
 	}
 

pkg/sql/vecindex/cspann/index.go

@@ -79,7 +79,7 @@ func TestIndex(t *testing.T) {
 			case "delete":
 				return state.Delete(d)
 
-			case "force-split", "force-merge":
+			case "force-split-or-merge":
 				return state.ForceSplitOrMerge(d)
 
 			case "recall":
@@ -475,10 +475,8 @@ func (s *testState) ForceSplitOrMerge(d *datadriven.TestData) string {
 		}
 	}
 
-	if d.Cmd == "force-split" {
-		s.Index.ForceSplit(s.Ctx, treeKey, parentPartitionKey, partitionKey)
-	} else {
-		s.Index.ForceMerge(s.Ctx, treeKey, parentPartitionKey, partitionKey)
+	if d.Cmd == "force-split-or-merge" {
+		s.Index.ForceSplitOrMerge(s.Ctx, treeKey, parentPartitionKey, partitionKey)
 	}
 
 	// Ensure the fixup runs.