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.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/partition.go

@@ -154,7 +154,7 @@ func (p *Partition) Search(
 	w *workspace.T, partitionKey PartitionKey, queryVector vector.T, searchSet *SearchSet,
 ) int {
 	count := p.Count()
-	tempFloats := w.AllocFloats(count * 2)
+	tempFloats := w.AllocFloats(count * 3)
 	defer w.FreeFloats(tempFloats)
 
 	// Estimate distances of the data vectors from the query vector.
@@ -163,6 +163,11 @@ func (p *Partition) Search(
 	p.quantizer.EstimateDistances(
 		w, p.quantizedSet, queryVector, tempDistances, tempErrorBounds)
 
+	tempCentroidDistances := tempFloats[count*2 : count*3]
+	if searchSet.IncludeCentroidDistances {
+		p.quantizer.GetCentroidDistances(p.quantizedSet, tempCentroidDistances, true /* spherical */)
+	}
+
 	// Add candidates to the search set, which is responsible for retaining the
 	// top-k results.
 	for i := range tempDistances {
@@ -173,6 +178,9 @@ func (p *Partition) Search(
 			ChildKey:           p.childKeys[i],
 			ValueBytes:         p.valueBytes[i],
 		}
+		if searchSet.IncludeCentroidDistances {
+			searchSet.tempResult.CentroidDistance = tempCentroidDistances[i]
+		}
 		searchSet.Add(&searchSet.tempResult)
 	}
 

pkg/sql/vecindex/cspann/partition_test.go

@@ -44,10 +44,11 @@ func TestPartition(t *testing.T) {
 	valueBytes20b := ValueBytes{11, 12}
 
 	var workspace workspace.T
-	quantizer := quantize.NewUnQuantizer(2, vecpb.L2SquaredDistance)
+	unquantizer := quantize.NewUnQuantizer(2, vecpb.L2SquaredDistance)
+	rabitq := quantize.NewRaBitQuantizer(2, 42, vecpb.InnerProductDistance)
 
 	// newTestPartition creates a partition with 2 vectors.
-	newTestPartition := func() *Partition {
+	newTestPartition := func(quantizer quantize.Quantizer) *Partition {
 		vectors := vector.MakeSet(2)
 		vectors.Add(vec10)
 		vectors.Add(vec20)
@@ -69,7 +70,7 @@ func TestPartition(t *testing.T) {
 
 	t.Run("test Init", func(t *testing.T) {
 		// Validate that Init sets same values.
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		var partition2 Partition
 		partition2.Init(
 			*partition.Metadata(),
@@ -82,7 +83,7 @@ func TestPartition(t *testing.T) {
 	})
 
 	t.Run("test Clone", func(t *testing.T) {
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		cloned := partition.Clone()
 		require.Equal(t, partition, cloned)
 
@@ -101,7 +102,7 @@ func TestPartition(t *testing.T) {
 	})
 
 	t.Run("test Add", func(t *testing.T) {
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		require.True(t, partition.Add(&workspace, vec40, childKey40, valueBytes40, true /* overwrite */))
 		require.Equal(t, 4, partition.Count())
 		require.Equal(t, []ChildKey{childKey10, childKey20, childKey30, childKey40}, partition.ChildKeys())
@@ -111,14 +112,14 @@ func TestPartition(t *testing.T) {
 		checkPartitionMetadata(t, partition.Metadata(), Level(1), vector.T{4, 3.33})
 
 		// Add vector with duplicate key and overwrite=false. Expect no-op.
-		partition = newTestPartition()
+		partition = newTestPartition(unquantizer)
 		require.False(t, partition.Add(&workspace, vec20b, childKey20, valueBytes20b, false /* overwrite */))
 		require.Equal(t, 3, partition.Count())
 		require.Equal(t, []ValueBytes{valueBytes10, valueBytes20, valueBytes30}, partition.ValueBytes())
 
 		// Add vector with duplicate key and overwrite=true. Expect value to be
 		// updated.
-		partition = newTestPartition()
+		partition = newTestPartition(unquantizer)
 		require.False(t, partition.Add(&workspace, vec20b, childKey20, valueBytes20b, true /* overwrite */))
 		require.Equal(t, 3, partition.Count())
 		require.Equal(t, []ChildKey{childKey10, childKey30, childKey20}, partition.ChildKeys())
@@ -129,7 +130,7 @@ func TestPartition(t *testing.T) {
 	t.Run("test AddSet", func(t *testing.T) {
 		// Create empty partition.
 		metadata := PartitionMetadata{Level: 1, Centroid: vector.T{4, 3}}
-		partition := CreateEmptyPartition(quantizer, metadata)
+		partition := CreateEmptyPartition(unquantizer, metadata)
 
 		// Add empty set.
 		vectors := vector.MakeSet(2)
@@ -188,7 +189,7 @@ func TestPartition(t *testing.T) {
 	t.Run("test Search", func(t *testing.T) {
 		// Search empty partition.
 		metadata := PartitionMetadata{Level: LeafLevel, Centroid: vector.T{4, 3}}
-		partition := CreateEmptyPartition(quantizer, metadata)
+		partition := CreateEmptyPartition(unquantizer, metadata)
 		require.Equal(t, Level(1), partition.Level())
 
 		searchSet := SearchSet{MaxResults: 1}
@@ -198,7 +199,7 @@ func TestPartition(t *testing.T) {
 		require.Equal(t, SearchResults(nil), results)
 
 		// Search partition with 5 vectors.
-		partition = newTestPartition()
+		partition = newTestPartition(unquantizer)
 		vectors := vector.MakeSet(2)
 		vectors.Add(vec40)
 		vectors.Add(vec50)
@@ -219,8 +220,24 @@ func TestPartition(t *testing.T) {
 		require.Equal(t, SearchResults{result1, result2, result3}, results)
 	})
 
+	t.Run("test Search with IncludeCentroidDistances", func(t *testing.T) {
+		// Search partition with 3 vectors.
+		partition := newTestPartition(rabitq)
+
+		searchSet := SearchSet{MaxResults: 2, IncludeCentroidDistances: true}
+		_ = partition.Search(&workspace, RootKey, vector.T{1, 1}, &searchSet)
+		result1 := SearchResult{
+			QueryDistance: -11.52, ErrorBound: 8.96, CentroidDistance: -8.45, ParentPartitionKey: 1,
+			ChildKey: childKey30, ValueBytes: valueBytes30}
+		result2 := SearchResult{
+			QueryDistance: -6.1, ErrorBound: 4.48, CentroidDistance: -5.12, ParentPartitionKey: 1,
+			ChildKey: childKey20, ValueBytes: valueBytes20}
+		results := roundResults(searchSet.PopResults(), 2)
+		require.Equal(t, SearchResults{result1, result2}, results)
+	})
+
 	t.Run("test ReplaceWithLast", func(t *testing.T) {
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		partition.ReplaceWithLast(0)
 		require.Equal(t, 2, partition.Count())
 		require.Equal(t, []ChildKey{childKey30, childKey20}, partition.ChildKeys())
@@ -237,7 +254,7 @@ func TestPartition(t *testing.T) {
 	})
 
 	t.Run("test ReplaceWithLastByKey and Find", func(t *testing.T) {
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		require.Equal(t, 0, partition.Find(childKey10))
 		require.Equal(t, 2, partition.Find(childKey30))
 		require.True(t, partition.ReplaceWithLastByKey(childKey10))
@@ -252,7 +269,7 @@ func TestPartition(t *testing.T) {
 	})
 
 	t.Run("test Clear", func(t *testing.T) {
-		partition := newTestPartition()
+		partition := newTestPartition(unquantizer)
 		require.Equal(t, 3, partition.Clear())
 		require.Equal(t, 0, partition.Count())
 		require.Equal(t, []ChildKey{}, partition.ChildKeys())
@@ -273,6 +290,7 @@ func roundResults(results SearchResults, prec int) SearchResults {
 		result := &results[i]
 		result.QueryDistance = float32(scalar.Round(float64(result.QueryDistance), prec))
 		result.ErrorBound = float32(scalar.Round(float64(result.ErrorBound), prec))
+		result.CentroidDistance = float32(scalar.Round(float64(result.CentroidDistance), prec))
 		result.Vector = testutils.RoundFloats(result.Vector, prec)
 	}
 	return results

pkg/sql/vecindex/cspann/search_set.go

@@ -26,6 +26,12 @@ type SearchResult struct {
 	// ErrorBound captures the uncertainty of the distance estimate, which is
 	// highly likely to fall within QueryDistance ± ErrorBound.
 	ErrorBound float32
+	// CentroidDistance is the exact distance of the result vector from its
+	// centroid, according to the distance metric in use (e.g. L2Squared or
+	// Cosine).
+	// NOTE: This is only returned if the SearchSet.IncludeCentroidDistances is
+	// set to true.
+	CentroidDistance float32
 	// ParentPartitionKey is the key of the parent of the partition that contains
 	// the data vector.
 	ParentPartitionKey PartitionKey
@@ -182,6 +188,15 @@ type SearchSet struct {
 	// matching primary key.
 	MatchKey KeyBytes
 
+	// ExcludedPartitions specifies which partitions to skip during search.
+	// Vectors in any of these partitions will not be added to the set.
+	ExcludedPartitions []PartitionKey
+
+	// IncludeCentroidDistances indicates that search results need to have their
+	// CentroidDistance field set. This records the vector's distance from the
+	// centroid of its partition.
+	IncludeCentroidDistances bool
+
 	// Stats tracks useful information about the search, such as how many vectors
 	// and partitions were scanned.
 	Stats SearchStats
@@ -210,7 +225,7 @@ type SearchSet struct {
 func (ss *SearchSet) Init() {
 	ss.deDuper.Clear()
 	*ss = SearchSet{
-		deDuper: ss.deDuper,
+		deDuper:    ss.deDuper,
 		candidates: ss.candidates[:0],
 	}
 }
@@ -222,7 +237,8 @@ func (ss *SearchSet) Count() int {
 	return len(ss.candidates)
 }
 
-// Clear removes all candidates from the set.
+// Clear removes all candidates from the set, but does not otherwise disturb
+// other settings.
 func (ss *SearchSet) Clear() {
 	ss.candidates = ss.candidates[:0]
 	ss.deDuper.Clear()
@@ -254,6 +270,13 @@ func (ss *SearchSet) Add(candidate *SearchResult) {
 		return
 	}
 
+	// Skip vectors in excluded partitions.
+	if ss.ExcludedPartitions != nil {
+		if slices.Contains(ss.ExcludedPartitions, candidate.ParentPartitionKey) {
+			return
+		}
+	}
+
 	if ss.candidates == nil {
 		// Pre-allocate some capacity for candidates.
 		ss.candidates = make(searchResultHeap, 0, 16)
@@ -277,7 +300,7 @@ func (ss *SearchSet) AddSet(searchSet *SearchSet) {
 		return
 	}
 	ss.candidates = slices.Grow(ss.candidates, len(searchSet.candidates))
-	if ss.MatchKey != nil {
+	if ss.MatchKey != nil || ss.ExcludedPartitions != nil {
 		// Add each candidate individually in order to check the match key.
 		ss.AddAll(SearchResults(searchSet.candidates))
 	} else {

pkg/sql/vecindex/cspann/search_set_test.go

@@ -266,4 +266,19 @@ func TestSearchSet(t *testing.T) {
 		require.Equal(t, []float64{3, 3, 1, 4, 3, 0.5, 5, 4}, searchSet.FindBestDistances(distances[:8]))
 		require.Equal(t, []float64{3, 6, 1, 4, 3, 6, 5, 4, 3, 0.5}, searchSet.FindBestDistances(distances[:12]))
 	})
+
+	t.Run("test ExcludedPartitions", func(t *testing.T) {
+		result9 := SearchResult{
+			QueryDistance: 9, ErrorBound: 1, ParentPartitionKey: 500, ChildKey: ChildKey{KeyBytes: []byte{90}}}
+		searchSet := SearchSet{MaxResults: 9, ExcludedPartitions: []PartitionKey{100, 500, 800}}
+		searchSet.AddAll(SearchResults{
+			result1, result2, result3, result4, result1, result5, result6, result7, result1, result8, result9})
+		require.Equal(t, SearchResults{result3, result4, result7, result6, result2}, searchSet.PopResults())
+
+		set1 := SearchSet{MaxResults: 3, ExcludedPartitions: []PartitionKey{200}}
+		set2 := SearchSet{MaxResults: 3}
+		set2.AddAll(SearchResults{result1, result2, result3})
+		set1.AddSet(&set2)
+		require.Equal(t, SearchResults{result3, result1}, set1.PopResults())
+	})
 }