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.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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
 			}

pkg/sql/vecindex/cspann/fixup_split.go

@@ -97,6 +97,8 @@ type fixupWorker struct {
 	tempChildKey        [1]ChildKey
 	tempValueBytes      [1]ValueBytes
 	tempMetadataToGet   []PartitionMetadataToGet
+	tempIndexCtx        Context
+	tempPartitionKeys   [3]PartitionKey
 }
 
 // ewFixupWorker returns a new worker for the given processor.
@@ -240,24 +242,33 @@ func (fw *fixupWorker) getFullVectorsForPartition(
 		vectors = vector.MakeSet(fw.index.quantizer.GetDims())
 		vectors.AddUndefined(len(fw.tempVectorsWithKeys))
 		for i := range fw.tempVectorsWithKeys {
-			if partition.Level() == LeafLevel {
-				// Leaf vectors from the primary index need to be randomized and
-				// possibly normalized.
-				fw.index.TransformVector(fw.tempVectorsWithKeys[i].Vector, vectors.At(i))
-			} else {
-				copy(vectors.At(i), fw.tempVectorsWithKeys[i].Vector)
-
-				// Convert mean centroids into spherical centroids for the Cosine
-				// and InnerProduct distance metrics.
-				switch fw.index.quantizer.GetDistanceMetric() {
-				case vecpb.CosineDistance, vecpb.InnerProductDistance:
-					num32.Normalize(vectors.At(i))
-				}
-			}
+			fw.transformFullVector(partition.Level(), fw.tempVectorsWithKeys[i].Vector, vectors.At(i))
 		}
 
 		return nil
 	})
 
 	return vectors, err
 }
+
+// transformFullVector ensures that the full vector fetched from a partition at
+// the given level has been properly randomized and normalized. It copies the
+// randomized, normalized vector into "randomized", which must be allocated by
+// the caller with the same length as the input vector.
+func (fw *fixupWorker) transformFullVector(level Level, vec, randomized vector.T) {
+	if level == LeafLevel {
+		// Leaf vectors from the primary index need to be randomized and possibly
+		// normalized.
+		fw.index.TransformVector(vec, randomized)
+	} else {
+		// This is an interior level, which means the vector is a partition
+		// centroid that's already normalized. However, it's a mean centroid, and
+		// needs to be converted into a spherical centroid for the Cosine and
+		// InnerProduct distance metrics.
+		copy(randomized, vec)
+		switch fw.index.quantizer.GetDistanceMetric() {
+		case vecpb.CosineDistance, vecpb.InnerProductDistance:
+			num32.Normalize(randomized)
+		}
+	}
+}

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -474,7 +474,7 @@ func (vi *Index) Search(
 	searchSet *SearchSet,
 	options SearchOptions,
 ) error {
-	vi.setupContext(idxCtx, treeKey, vec, options, LeafLevel)
+	vi.setupContext(idxCtx, treeKey, options, LeafLevel, vec, false /* transformed */)
 	return vi.searchHelper(ctx, idxCtx, searchSet)
 }
 
@@ -595,11 +595,11 @@ func (vi *Index) ForceMerge(
 func (vi *Index) setupInsertContext(idxCtx *Context, treeKey TreeKey, vec vector.T) {
 	// Perform the search using quantized vectors rather than full vectors (i.e.
 	// skip reranking).
-	vi.setupContext(idxCtx, treeKey, vec, SearchOptions{
+	vi.setupContext(idxCtx, treeKey, SearchOptions{
 		BaseBeamSize: vi.options.BaseBeamSize,
 		SkipRerank:   true,
 		UpdateStats:  !vi.options.DisableAdaptiveSearch,
-	}, SecondLevel)
+	}, SecondLevel, vec, false /* transformed */)
 	idxCtx.forInsert = true
 }
 
@@ -608,27 +608,39 @@ func (vi *Index) setupDeleteContext(idxCtx *Context, treeKey TreeKey, vec vector
 	// Perform the search using quantized vectors rather than full vectors (i.e.
 	// skip reranking). Use a larger beam size to make it more likely that we'll
 	// find the vector to delete.
-	vi.setupContext(idxCtx, treeKey, vec, SearchOptions{
+	vi.setupContext(idxCtx, treeKey, SearchOptions{
 		BaseBeamSize: vi.options.BaseBeamSize * 2,
 		SkipRerank:   true,
 		UpdateStats:  !vi.options.DisableAdaptiveSearch,
-	}, LeafLevel)
+	}, LeafLevel, vec, false /* transformed */)
 	idxCtx.forDelete = true
 }
 
-// setupContext sets up the given context as an operation is beginning.
+// setupContext sets up the given context as an operation is beginning. If
+// "randomized" is false, then the given vector is expected to be an original,
+// unrandomized vector that was provided by the user. If "transformed" is true,
+// then the vector is expected to already be randomized and normalized.
 func (vi *Index) setupContext(
-	idxCtx *Context, treeKey TreeKey, vec vector.T, options SearchOptions, level Level,
+	idxCtx *Context,
+	treeKey TreeKey,
+	options SearchOptions,
+	level Level,
+	vec vector.T,
+	transformed bool,
 ) {
 	idxCtx.treeKey = treeKey
 	idxCtx.level = level
-	idxCtx.query.Init(vi.quantizer.GetDistanceMetric(), vec, &vi.rot)
 	idxCtx.forInsert = false
 	idxCtx.forDelete = false
 	idxCtx.options = options
 	if idxCtx.options.BaseBeamSize == 0 {
 		idxCtx.options.BaseBeamSize = vi.options.BaseBeamSize
 	}
+	if transformed {
+		idxCtx.query.InitTransformed(vi.quantizer.GetDistanceMetric(), vec, &vi.rot)
+	} else {
+		idxCtx.query.InitOriginal(vi.quantizer.GetDistanceMetric(), vec, &vi.rot)
+	}
 }
 
 // updateFunc is called by searchForUpdateHelper when it has a candidate
@@ -920,6 +932,25 @@ func (vi *Index) findExactDistances(
 		return candidates, nil
 	}
 
+	var err error
+	candidates, err = vi.getFullVectors(ctx, idxCtx, candidates)
+	if err != nil {
+		return nil, errors.Wrapf(err, "getting full vectors to find exact distances")
+	}
+
+	// Compute exact distance between query vector and the data vectors.
+	idxCtx.query.ComputeExactDistances(idxCtx.level, candidates)
+
+	return candidates, nil
+}
+
+// getFullVectors fetches the full-size vectors for the given search candidates.
+// These can be either leaf vectors fetched from the primary index or interior
+// partition centroids fetched from the index. Fixups are enqueued for any
+// vectors found to be "dangling".
+func (vi *Index) getFullVectors(
+	ctx context.Context, idxCtx *Context, candidates []SearchResult,
+) ([]SearchResult, error) {
 	// Prepare vector references.
 	idxCtx.tempVectorsWithKeys = utils.EnsureSliceLen(idxCtx.tempVectorsWithKeys, len(candidates))
 	for i := range candidates {
@@ -956,9 +987,6 @@ func (vi *Index) findExactDistances(
 		}
 	}
 
-	// Compute exact distance between query vector and the data vectors.
-	idxCtx.query.ComputeExactDistances(idxCtx.level, candidates)
-
 	return candidates, nil
 }
 

pkg/sql/vecindex/cspann/index.go

@@ -446,6 +446,8 @@ func (s *testState) ForceSplitOrMerge(d *datadriven.TestData) string {
 
 		case "steps":
 			steps = s.parseInt(arg)
+			// Always discard any fixups triggered by the split or merge fixup.
+			s.DiscardFixups = true
 		}
 	}
 

pkg/sql/vecindex/cspann/index_test.go

@@ -28,23 +28,35 @@ type queryComparer struct {
 	randomized vector.T
 }
 
-// Init sets the query vector and prepares the comparer for use.
-func (c *queryComparer) Init(
-	distanceMetric vecpb.DistanceMetric, queryVector vector.T, rot *RandomOrthoTransformer,
+// InitOriginal sets the original query vector and prepares the comparer for
+// use.
+func (c *queryComparer) InitOriginal(
+	distanceMetric vecpb.DistanceMetric, original vector.T, rot *RandomOrthoTransformer,
 ) {
 	c.distanceMetric = distanceMetric
-	c.original = queryVector
+	c.original = original
 
 	// Randomize the original query vector.
-	c.randomized = utils.EnsureSliceLen(c.randomized, len(queryVector))
-	c.randomized = rot.RandomizeVector(queryVector, c.randomized)
+	c.randomized = utils.EnsureSliceLen(c.randomized, len(original))
+	c.randomized = rot.RandomizeVector(original, c.randomized)
 
 	// If using cosine distance, also normalize the query vector.
 	if c.distanceMetric == vecpb.CosineDistance {
 		num32.Normalize(c.randomized)
 	}
 }
 
+// InitRandomized sets the transformed query vector in cases where the original
+// query vector is not available, such as when the vector is an interior
+// partition centroid. It is expected to already be randomized and normalized.
+func (c *queryComparer) InitTransformed(
+	distanceMetric vecpb.DistanceMetric, randomized vector.T, rot *RandomOrthoTransformer,
+) {
+	c.distanceMetric = distanceMetric
+	c.original = nil
+	c.randomized = randomized
+}
+
 // Randomized returns the query vector after it has been randomized and
 // normalized as needed.
 func (c *queryComparer) Randomized() vector.T {

pkg/sql/vecindex/cspann/query_comparer.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecpb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/num32"
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 	"github.com/stretchr/testify/require"
 )
@@ -158,16 +159,49 @@ func TestQueryComparer(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			// Setup queryComparer.
 			var rot RandomOrthoTransformer
-			rot.Init(vecpb.RotNone, len(tc.queryVector), 42)
+			rot.Init(vecpb.RotGivens, len(tc.queryVector), 42)
 
 			var comparer queryComparer
-			comparer.Init(tc.metric, tc.queryVector, &rot)
+			comparer.InitOriginal(tc.metric, tc.queryVector, &rot)
 
-			// Make a copy of candidates to avoid modifying test data.
+			// Make a copy of the candidates.
 			candidates := make([]SearchResult, len(tc.candidates))
 			copy(candidates, tc.candidates)
 
-			// Test the main method.
+			// Randomize candidates from an interior level.
+			if tc.level != LeafLevel {
+				for i := range tc.candidates {
+					rot.RandomizeVector(tc.candidates[i].Vector, candidates[i].Vector)
+				}
+			}
+
+			// Test ComputeExactDistances.
+			comparer.ComputeExactDistances(tc.level, candidates)
+			require.Len(t, candidates, len(tc.expected), "number of candidates should be preserved")
+
+			for i, expected := range tc.expected {
+				require.InDelta(t, expected, candidates[i].QueryDistance, 1e-5,
+					"distance mismatch for candidate %d", i)
+
+				// Error bound should always be 0 for exact distances.
+				require.Equal(t, float32(0), candidates[i].ErrorBound,
+					"error bound should be 0 for exact distances")
+			}
+
+			// Test InitRandomized for interior levels.
+			if tc.level == LeafLevel {
+				return
+			}
+
+			// Transform the query vector.
+			queryVector := make([]float32, len(tc.queryVector))
+			rot.RandomizeVector(tc.queryVector, queryVector)
+			if tc.metric == vecpb.CosineDistance {
+				num32.Normalize(queryVector)
+			}
+			comparer.InitTransformed(tc.metric, queryVector, &rot)
+
+			// Test ComputeExactDistances.
 			comparer.ComputeExactDistances(tc.level, candidates)
 			require.Len(t, candidates, len(tc.expected), "number of candidates should be preserved")
 

pkg/sql/vecindex/cspann/query_comparer_test.go

@@ -8,6 +8,7 @@ package cspann
 import (
 	"context"
 	"math"
+	"slices"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/utils"
 	"github.com/cockroachdb/errors"
@@ -172,6 +173,8 @@ type levelSearcher struct {
 	parent *levelSearcher
 	// stats points to search stats that should be updated as the search runs.
 	stats *SearchStats
+	// excludedPartitions specifies which partitions to skip during search.
+	excludedPartitions []PartitionKey
 	// level is the K-means tree level being searched. This is undefined for the
 	// root level until SearchRoot is called.
 	level Level
@@ -244,7 +247,11 @@ func (s *levelSearcher) Init(
 			}
 			s.searchSet.MaxResults = maxResults
 			s.searchSet.MaxExtraResults = maxExtraResults
+
+			// Set additional fields that only apply to the last level.
 			s.searchSet.MatchKey = searchSet.MatchKey
+			s.searchSet.IncludeCentroidDistances = searchSet.IncludeCentroidDistances
+			s.excludedPartitions = searchSet.ExcludedPartitions
 		}
 	}
 }
@@ -298,12 +305,27 @@ func (s *levelSearcher) NextBatch(ctx context.Context) (ok bool, err error) {
 	// overlap with the previous batch, in terms of ordering and duplicates.
 	s.searchSet.Clear()
 
+	// filterPartitions filters parent results so that we don't even attempt to
+	// search excluded partitions.
+	filterPartitions := func(results []SearchResult) []SearchResult {
+		if s.excludedPartitions == nil {
+			return results
+		}
+		for i := 0; i < len(results); i++ {
+			if slices.Contains(s.excludedPartitions, results[i].ChildKey.PartitionKey) {
+				results = utils.ReplaceWithLast(results, i)
+				i--
+			}
+		}
+		return results
+	}
+
 	if firstBatch {
 		ok, err := s.parent.NextBatch(ctx)
 		if err != nil || !ok {
 			return ok, err
 		}
-		s.parentResults = s.parent.SearchSet().PopResults()
+		s.parentResults = filterPartitions(s.parent.SearchSet().PopResults())
 	} else if len(s.parentResults) < s.beamSize {
 		// Get more results from parent to try and fill the beam size.
 		parentResults := s.parent.SearchSet().PopResults()
@@ -319,6 +341,7 @@ func (s *levelSearcher) NextBatch(ctx context.Context) (ok bool, err error) {
 			}
 			parentResults = s.parent.SearchSet().PopResults()
 		}
+		parentResults = filterPartitions(parentResults)
 		if len(s.parentResults) == 0 {
 			s.parentResults = parentResults
 		} else {