vecindex: return minimum number of results from Search

When vector deletion is enabled in TestVecIndexConcurrency, it occasionally
fails under stress with zero results. This is because searching with a LIMIT
of one was returning only 2 results. If both of those results point to deleted
primary key rows, then zero results get returned. The fix is to return at
least 10 candidate vectors from searches. While all 10 vectors could be deleted,
this at least makes this scenario much less common. The longer-term fix is to
switch to an iterator execution model, where we continue to "pull" additional
vectors until we get the desired number of results.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/index.go

@@ -8,6 +8,7 @@ package cspann
 import (
 	"bytes"
 	"context"
+	"math"
 	"math/rand"
 	"runtime"
 	"strconv"
@@ -28,13 +29,33 @@ import (
 // of search results that will be reranked with the original full-size vectors.
 const RerankMultiplier = 10
 
-// DeletedMultiplier increases the number of results that will be reranked, in
+// DeletedMinCount sets a minimum number of results that will be reranked, in
 // order to account for vectors that may have been deleted in the primary index.
+const DeletedMinCount = 10
+
+// DeletedMultiplier increases the number of results that will be reranked by
+// this factor, in order to account for vectors that may have been deleted in
+// the primary index.
 const DeletedMultiplier = 1.2
 
 // MaxQualitySamples specifies the max value of the QualitySamples index option.
 const MaxQualitySamples = 32
 
+// IncreaseRerankResults returns good values for maxResults and maxExtraResults
+// that have a high probability of returning the desired number of results, even
+// when there are deleted results. Deleted results will be filtered out by the
+// rerank process, so we need to make sure there are additional results that can
+// be returned instead.
+//
+// TODO(andyk): Switch the index to use a search iterator so the caller can keep
+// requesting further results rather than guessing at how many additional
+// results might be needed.
+func IncreaseRerankResults(desiredMaxResults int) (maxResults, maxExtraResults int) {
+	maxResults = max(int(math.Ceil(float64(desiredMaxResults)*DeletedMultiplier)), DeletedMinCount)
+	maxExtraResults = desiredMaxResults * RerankMultiplier
+	return maxResults, maxExtraResults
+}
+
 // IndexOptions specifies options that control how the index will be built, as
 // well as default options for how it will be searched. A given search operation
 // can specify SearchOptions to override the default behavior.

pkg/sql/vecindex/cspann/searcher.go

@@ -120,8 +120,10 @@ func (s *searcher) Next(ctx context.Context) (ok bool, err error) {
 				maxResults = s.searchSet.MaxResults
 				maxExtraResults = s.searchSet.MaxExtraResults
 				if !s.idxCtx.options.SkipRerank {
-					maxResults = int(math.Ceil(float64(maxResults) * DeletedMultiplier))
-					maxExtraResults = maxResults * RerankMultiplier
+					maxResults, maxExtraResults = IncreaseRerankResults(maxResults)
+					if s.searchSet.MaxExtraResults > maxExtraResults {
+						maxExtraResults = s.searchSet.MaxExtraResults
+					}
 				}
 				if s.idxCtx.level != LeafLevel && s.idxCtx.options.UpdateStats {
 					maxResults = max(maxResults, s.idx.options.QualitySamples)

pkg/sql/vecindex/cspann/testdata/read-only.ddt

@@ -37,7 +37,7 @@ vec4: 2 (centroid=1.25)
 vec10: 4 (centroid=1.75)
 vec5: 4 (centroid=2.02)
 vec8: 4 (centroid=2.25)
-13 leaf vectors, 16 vectors, 5 full vectors, 4 partitions
+13 leaf vectors, 16 vectors, 13 full vectors, 4 partitions
 
 format-tree
 ----
@@ -71,7 +71,7 @@ search max-results=3 beam-size=3
 vec9: 18 (centroid=4.95)
 vec11: 25 (centroid=2.36)
 vec5: 26 (centroid=2.02)
-13 leaf vectors, 16 vectors, 10 full vectors, 4 partitions
+13 leaf vectors, 16 vectors, 11 full vectors, 4 partitions
 
 format-tree
 ----

pkg/sql/vecindex/cspann/testdata/search-features.ddt

@@ -14,7 +14,7 @@ CV stats:
 search max-results=1 use-feature=5000 beam-size=1
 ----
 vec771: 0.5624 (centroid=0.63)
-23 leaf vectors, 43 vectors, 2 full vectors, 4 partitions
+23 leaf vectors, 43 vectors, 12 full vectors, 4 partitions
 
 # Search for additional results.
 search max-results=6 use-feature=5000 beam-size=1
@@ -36,7 +36,7 @@ vec640: 0.6525 (centroid=0.52)
 vec329: 0.6871 (centroid=0.52)
 vec95: 0.7008 (centroid=0.61)
 vec386: 0.7301 (centroid=0.61)
-85 leaf vectors, 141 vectors, 13 full vectors, 13 partitions
+85 leaf vectors, 141 vectors, 18 full vectors, 13 partitions
 
 # Turn off re-ranking, which results in increased inaccuracy.
 search max-results=6 use-feature=5000 beam-size=4 skip-rerank

pkg/sql/vecindex/cspann/testdata/search.ddt

@@ -217,3 +217,64 @@ vec2: 28.1958 ±10.12 (centroid=1.58)
 vec3: 45.6359 ±15.23 (centroid=1.41)
 vec5: 74.3641 ±15.23 (centroid=1.41)
 4 leaf vectors, 7 vectors, 0 full vectors, 3 partitions
+
+# ----------------------------------------------------------------------
+# Search tree with many deleted vectors.
+# ----------------------------------------------------------------------
+new-index min-partition-size=1 max-partition-size=4 beam-size=3
+vec1: (0, 0)
+vec2: (100, 100)
+vec3: (200, 200)
+vec4: (300, 300)
+vec5: (400, 400)
+vec6: (500, 500)
+----
+• 1 (0, 0)
+│
+├───• 5 (450, 450)
+│   │
+│   ├───• vec5 (400, 400)
+│   └───• vec6 (500, 500)
+│
+├───• 3 (50, 50)
+│   │
+│   ├───• vec1 (0, 0)
+│   └───• vec2 (100, 100)
+│
+└───• 4 (250, 250)
+    │
+    ├───• vec3 (200, 200)
+    └───• vec4 (300, 300)
+
+# Delete all but one vector.
+delete not-found
+vec1
+vec2
+vec3
+vec4
+vec5
+----
+• 1 (0, 0)
+│
+├───• 5 (450, 450)
+│   │
+│   ├───• vec5 (MISSING)
+│   └───• vec6 (500, 500)
+│
+├───• 3 (50, 50)
+│   │
+│   ├───• vec1 (MISSING)
+│   └───• vec2 (MISSING)
+│
+└───• 4 (250, 250)
+    │
+    ├───• vec3 (MISSING)
+    └───• vec4 (MISSING)
+
+# Search for the vector that's farthest from vec6, so it's last in the result
+# set. This ensures that DeletedMinCount is working as intended.
+search max-results=1
+(0, 0)
+----
+vec6: 500000 (centroid=70.71)
+6 leaf vectors, 9 vectors, 6 full vectors, 4 partitions

pkg/sql/vecindex/cspann/testdata/split-non-root-step.ddt

@@ -781,7 +781,7 @@ search beam-size=16
 (2, -4)
 ----
 vec5: 0 (centroid=1)
-7 leaf vectors, 17 vectors, 3 full vectors, 11 partitions
+7 leaf vectors, 17 vectors, 5 full vectors, 11 partitions
 
 format-tree
 ----
@@ -820,7 +820,7 @@ search beam-size=16
 (2, -4)
 ----
 vec5: 0 (centroid=1)
-5 leaf vectors, 17 vectors, 3 full vectors, 13 partitions
+5 leaf vectors, 17 vectors, 5 full vectors, 13 partitions
 
 format-tree
 ----
@@ -863,7 +863,7 @@ search beam-size=16
 (2, -4)
 ----
 vec5: 0 (centroid=1)
-5 leaf vectors, 19 vectors, 3 full vectors, 15 partitions
+5 leaf vectors, 19 vectors, 5 full vectors, 15 partitions
 
 format-tree
 ----
@@ -1042,7 +1042,7 @@ search
 vec2: 0 (centroid=0)
 vec6: 0 (centroid=0)
 vec7: 0 (centroid=0)
-5 leaf vectors, 12 vectors, 4 full vectors, 6 partitions
+5 leaf vectors, 12 vectors, 5 full vectors, 6 partitions
 
 format-tree
 ----

pkg/sql/vecindex/cspann/testdata/split-root-step.ddt

@@ -255,7 +255,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 13 vectors, 6 full vectors, 5 partitions
+9 leaf vectors, 13 vectors, 9 full vectors, 5 partitions
 
 # Move to the DrainingForSplit state, where sub-partition #6 has vectors.
 force-split partition-key=1 steps=4
@@ -308,7 +308,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 15 vectors, 6 full vectors, 7 partitions
+9 leaf vectors, 15 vectors, 9 full vectors, 7 partitions
 
 # Move to the point where partition #1 children have been cleared.
 force-split partition-key=1 steps=2
@@ -324,7 +324,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 13 vectors, 6 full vectors, 7 partitions
+9 leaf vectors, 13 vectors, 9 full vectors, 7 partitions
 
 # Move to the AddingLevel state, where root partition's level has increased.
 force-split partition-key=1 steps=1
@@ -340,7 +340,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 13 vectors, 6 full vectors, 7 partitions
+9 leaf vectors, 13 vectors, 9 full vectors, 7 partitions
 
 # Move to point where sub-partitions #6 and #7 have been added to the root.
 force-split partition-key=1 steps=4
@@ -381,7 +381,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 15 vectors, 6 full vectors, 7 partitions
+9 leaf vectors, 15 vectors, 9 full vectors, 7 partitions
 
 # Finish the split.
 force-split partition-key=1 steps=1
@@ -422,7 +422,7 @@ vec7: 5 (centroid=1.67)
 vec8: 5 (centroid=2.24)
 vec3: 26 (centroid=1.05)
 vec4: 26 (centroid=1.05)
-9 leaf vectors, 15 vectors, 6 full vectors, 7 partitions
+9 leaf vectors, 15 vectors, 9 full vectors, 7 partitions
 
 # ----------------------------------------------------------------------
 # Insert into the tree when the root is in splitting states.

pkg/sql/vecindex/cspann/testdata/split.ddt

@@ -147,7 +147,7 @@ search max-results=3 beam-size=3
 vec7: 2 (centroid=2.12)
 vec10: 5 (centroid=0.71)
 vec12: 9 (centroid=0.71)
-6 leaf vectors, 14 vectors, 5 full vectors, 6 partitions
+6 leaf vectors, 14 vectors, 6 full vectors, 6 partitions
 
 # ----------------------------------------------------------------------
 # Test linking nearby vectors from other partitions.

pkg/sql/vecindex/cspann/testdata/tree-key.ddt

@@ -156,7 +156,7 @@ search max-results=1 tree=1
 (1, 2)
 ----
 vec6: 5 (centroid=2.36)
-6 leaf vectors, 8 vectors, 2 full vectors, 3 partitions
+6 leaf vectors, 8 vectors, 6 full vectors, 3 partitions
 
 # Vector should now be gone from the index.
 format-tree tree=1

pkg/sql/vecindex/searcher.go

@@ -7,7 +7,6 @@ package vecindex
 
 import (
 	"context"
-	"math"
 
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -48,8 +47,7 @@ func (s *Searcher) Init(idx *cspann.Index, txn *kv.Txn, baseBeamSize, maxResults
 		BaseBeamSize: baseBeamSize,
 		SkipRerank:   true,
 	}
-	s.searchSet.MaxResults = int(math.Ceil(float64(maxResults) * cspann.DeletedMultiplier))
-	s.searchSet.MaxExtraResults = s.searchSet.MaxResults * cspann.RerankMultiplier
+	s.searchSet.MaxResults, s.searchSet.MaxExtraResults = cspann.IncreaseRerankResults(maxResults)
 
 	// If the index is deterministic, then synchronously run the background worker
 	// to process any pending fixups.