vecindex: improve vector reranking

This commit exposes a new vector_search_rerank_multiplier session setting that
controls how many vectors will be reranked at the end of a search. Empirical
testing shows that for difficult datasets (e.g. Glove and CLIP), the number of
vectors that need to be reranked is proportional to:

  log2(search_beam_size) * log2(top-k-results)

We multiply this number by the vector_search_rerank_multiplier setting, which is
set to 50 by default (also derived empirically). This formula will rerank enough
vectors to allow us to achieve high accuracies while still providing a safeguard
that prevents runaway evaluation in edge cases.

Less difficult datasets need to rerank far fewer vectors, but it doesn't hurt
to have a limit that's too high, since RaBitQ error bounds let us avoid actually
evaluating unneeded vectors.

This commit also increases the max value of vector_search_beam_size from 512 to
2048, since some datasets need it.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

pkg/cmd/vecbench/mem_provider.go

@@ -165,7 +165,9 @@ func (m *MemProvider) Search(
 		// Search the store.
 		var idxCtx cspann.Context
 		idxCtx.Init(txn)
-		searchSet := cspann.SearchSet{MaxResults: memState.maxResults}
+		maxResults, maxExtraResults :=
+			cspann.IncreaseRerankResults(memState.beamSize, memState.maxResults, 50)
+		searchSet := cspann.SearchSet{MaxResults: maxResults, MaxExtraResults: maxExtraResults}
 		searchOptions := cspann.SearchOptions{BaseBeamSize: memState.beamSize}
 		err = m.index.Search(ctx, &idxCtx, nil /* treeKey */, vec, &searchSet, searchOptions)
 		if err != nil {

pkg/sql/exec_util.go

@@ -4212,6 +4212,10 @@ func (m *sessionDataMutator) SetVectorSearchBeamSize(val int32) {
 	m.data.VectorSearchBeamSize = val
 }
 
+func (m *sessionDataMutator) SetVectorSearchRerankMultiplier(val int32) {
+	m.data.VectorSearchRerankMultiplier = val
+}
+
 func (m *sessionDataMutator) SetPropagateAdmissionHeaderToLeafTransactions(val bool) {
 	m.data.PropagateAdmissionHeaderToLeafTransactions = val
 }

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -4141,6 +4141,7 @@ use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
+vector_search_rerank_multiplier                                  50
 xmloption                                                        content
 
 # information_schema can be used with the anonymous database.

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -3146,6 +3146,7 @@ use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on                  NULL      NULL        NULL        string
 variable_inequality_lookup_join_enabled                          on                  NULL      NULL        NULL        string
 vector_search_beam_size                                          32                  NULL      NULL        NULL        string
+vector_search_rerank_multiplier                                  50                  NULL      NULL        NULL        string
 vectorize                                                        on                  NULL      NULL        NULL        string
 xmloption                                                        content             NULL      NULL        NULL        string
 
@@ -3382,6 +3383,7 @@ use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on                  NULL  user     NULL      on                  on
 variable_inequality_lookup_join_enabled                          on                  NULL  user     NULL      on                  on
 vector_search_beam_size                                          32                  NULL  user     NULL      32                  32
+vector_search_rerank_multiplier                                  50                  NULL  user     NULL      50                  50
 vectorize                                                        on                  NULL  user     NULL      on                  on
 xmloption                                                        content             NULL  user     NULL      content             content
 
@@ -3610,6 +3612,7 @@ use_pre_25_2_variadic_builtins                                   NULL    NULL
 use_proc_txn_control_extended_protocol_fix                       NULL    NULL     NULL     NULL        NULL
 variable_inequality_lookup_join_enabled                          NULL    NULL     NULL     NULL        NULL
 vector_search_beam_size                                          NULL    NULL     NULL     NULL        NULL
+vector_search_rerank_multiplier                                  NULL    NULL     NULL     NULL        NULL
 vectorize                                                        NULL    NULL     NULL     NULL        NULL
 xmloption                                                        NULL    NULL     NULL     NULL        NULL
 

pkg/sql/logictest/testdata/logic_test/show_source

@@ -250,6 +250,7 @@ use_pre_25_2_variadic_builtins                                   off
 use_proc_txn_control_extended_protocol_fix                       on
 variable_inequality_lookup_join_enabled                          on
 vector_search_beam_size                                          32
+vector_search_rerank_multiplier                                  50
 vectorize                                                        on
 xmloption                                                        content
 

pkg/sql/logictest/testdata/logic_test/vector_index

@@ -843,11 +843,31 @@ SHOW vector_search_beam_size;
 ----
 8
 
-statement error vector_search_beam_size cannot be less than 1 or greater than 512
+statement error vector_search_beam_size cannot be less than 1 or greater than 2048
 SET vector_search_beam_size=0
 
-statement error vector_search_beam_size cannot be less than 1 or greater than 512
-SET vector_search_beam_size=513
+statement error vector_search_beam_size cannot be less than 1 or greater than 2048
+SET vector_search_beam_size=2049
+
+# Ensure that the vector_search_rerank_multiplier session setting is settable.
+query T
+SHOW vector_search_rerank_multiplier;
+----
+50
+
+statement ok
+SET vector_search_rerank_multiplier=100
+
+query T
+SHOW vector_search_rerank_multiplier;
+----
+100
+
+statement error vector_search_rerank_multiplier cannot be less than 0 or greater than 100
+SET vector_search_rerank_multiplier=-1
+
+statement error vector_search_rerank_multiplier cannot be less than 0 or greater than 100
+SET vector_search_rerank_multiplier=101
 
 subtest end
 

pkg/sql/rowexec/vector_search.go

@@ -61,7 +61,9 @@ func newVectorSearchProcessor(
 	}
 	searchBeamSize := int(flowCtx.EvalCtx.SessionData().VectorSearchBeamSize)
 	maxResults := int(v.targetCount)
-	v.searcher.Init(flowCtx.EvalCtx, idx, flowCtx.Txn, &spec.GetFullVectorsFetchSpec, searchBeamSize, maxResults)
+	rerankMultiplier := int(flowCtx.EvalCtx.SessionData().VectorSearchRerankMultiplier)
+	v.searcher.Init(flowCtx.EvalCtx,
+		idx, flowCtx.Txn, &spec.GetFullVectorsFetchSpec, searchBeamSize, maxResults, rerankMultiplier)
 	colTypes := make([]*types.T, len(v.fetchSpec.FetchedColumns))
 	for i, col := range v.fetchSpec.FetchedColumns {
 		colTypes[i] = col.Type

pkg/sql/sessiondatapb/local_only_session_data.proto

@@ -694,6 +694,11 @@ message LocalOnlySessionData {
   bool enable_scrub_job = 176;
   // AllowViewWithSecurityInvokerClause indicates whether security invoker for views is enabled
   bool allow_view_with_security_invoker_clause = 177;
+  // VectorSearchRerankMultiplier controls how many of the initial search results
+  // can be reranked using exact distance calculations with the original
+  // full-size vectors. It acts as a multiplier on a base limit derived from the
+  // search beam size and the top-k results requested by the query.
+  int32 vector_search_rerank_multiplier = 178;
 
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //

pkg/sql/vars.go

@@ -4047,9 +4047,9 @@ var varGen = map[string]sessionVar{
 			if err != nil {
 				return err
 			}
-			if b < 1 || b > 512 {
+			if b < 1 || b > 2048 {
 				return pgerror.Newf(pgcode.InvalidParameterValue,
-					"vector_search_beam_size cannot be less than 1 or greater than 512")
+					"vector_search_beam_size cannot be less than 1 or greater than 2048")
 			}
 			m.SetVectorSearchBeamSize(int32(b))
 			return nil
@@ -4062,6 +4062,29 @@ var varGen = map[string]sessionVar{
 		},
 	},
 
+	// CockroachDB extension.
+	`vector_search_rerank_multiplier`: {
+		GetStringVal: makeIntGetStringValFn(`vector_search_rerank_multiplier`),
+		Set: func(_ context.Context, m sessionDataMutator, s string) error {
+			b, err := strconv.ParseInt(s, 10, 32)
+			if err != nil {
+				return err
+			}
+			if b < 0 || b > 100 {
+				return pgerror.Newf(pgcode.InvalidParameterValue,
+					"vector_search_rerank_multiplier cannot be less than 0 or greater than 100")
+			}
+			m.SetVectorSearchRerankMultiplier(int32(b))
+			return nil
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			return strconv.FormatInt(int64(evalCtx.SessionData().VectorSearchRerankMultiplier), 10), nil
+		},
+		GlobalDefault: func(sv *settings.Values) string {
+			return "50"
+		},
+	},
+
 	// CockroachDB extension.
 	`propagate_admission_header_to_leaf_transactions`: {
 		GetStringVal: makePostgresBoolGetStringValFn(`propagate_admission_header_to_leaf_transactions`),

pkg/sql/vecindex/cspann/index.go

@@ -25,10 +25,6 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
-// RerankMultiplier is multiplied by MaxResults to calculate the maximum number
-// of search results that will be reranked with the original full-size vectors.
-const RerankMultiplier = 10
-
 // 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
@@ -41,18 +37,30 @@ 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.
+// IncreaseRerankResults returns good values for maxResults and maxExtraResults.
+// Deleted results will be filtered out of the final results, so we need to make
+// sure there are additional results that can be returned instead. In addition,
+// quantization error can reduce the accuracy of results, so we need to return
+// extra results that can be reranked by exact distance calculations. Both the
+// search beam size and the top-k limit of results emperically have a
+// logarithmic relationship to the number of vectors that need to be reranked,
+// so use this formula to set a bound to the number of extra results:
+//
+// maxExtraResults =
+// log2(searchBeamSize) * log2(desiredMaxResults) * rerankMultiplier
+//
+// The rerank multiplier is a session setting that can be used to set a tighter
+// or looser bound.
 //
 // 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) {
+func IncreaseRerankResults(
+	searchBeamSize, desiredMaxResults, rerankMultiplier int,
+) (maxResults, maxExtraResults int) {
 	maxResults = max(int(math.Ceil(float64(desiredMaxResults)*DeletedMultiplier)), DeletedMinCount)
-	maxExtraResults = desiredMaxResults * RerankMultiplier
+	log := math.Log2(float64(max(searchBeamSize, 2))) * math.Log2(float64(max(desiredMaxResults, 2)))
+	maxExtraResults = int(log) * rerankMultiplier
 	return maxResults, maxExtraResults
 }