Merge #149847 #149858

149847: asim: more asim follow up  r=tbg a=wenyihu6

**asim: add replace = true to example_splitting**

Previously, we added a replace option to gen_load in data driven tests to
control whether the load generator is replaced or appended. The new default is
replace = false (append) which changed the original behaviour of replacing the
generator. This test expected the original behavior when re-running eval, but we
forgot to explicitly set replace = true. This commit fixes that.

Epic: none
Release note: none

---

**asim: add comment for BasicCluster.info**

This commit adds a comment to BasicCluster.info to highlight the inefficiency of
constructing the region node ratio when nodes_per_region is used with
gen_cluster just to reuse ClusterInfoWithDistribution.

Epic: none
Release note: none


149858: vecindex: improve vector reranking r=drewkimball a=andy-kimball

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

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

Author: 3 people

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/kv/kvserver/asim/gen/generator.go

@@ -232,6 +232,10 @@ func (bc BasicCluster) info() state.ClusterInfo {
 		return state.ClusterInfoWithStoreCount(bc.Nodes, bc.StoresPerNode)
 	}
 
+	// TODO(wenyihu6): we have the number of nodes and their localities already.
+	// We could construct ClusterInfo without a ratio calculation. We are doing
+	// this for now just to reuse ClusterInfoWithDistribution. But there may be
+	// rounding errors.
 	regionNodeWeights := make([]float64, len(bc.NodesPerRegion))
 	totalNodes := 0
 	for i, nodes := range bc.NodesPerRegion {

pkg/kv/kvserver/asim/tests/testdata/non_rand/example_splitting.txt

@@ -54,7 +54,7 @@ last store values: [s1=6] (stddev=0.00, mean=6.00, sum=6)
 # repeatedly - if the load split algorithm is good, we should expect the same
 # number of splits as the uniform workload. However, that is not the case as we
 # require more splits with a zipfian distribution.
-gen_load rate=10000 rw_ratio=1 access_skew=true
+gen_load rate=10000 rw_ratio=1 access_skew=true replace=true
 ----
 
 eval duration=5m samples=2 seed=42
@@ -63,24 +63,24 @@ OK
 
 plot stat=replicas sample=4
 ----
- 15.00 ┤                                ╭──────────────────────────────────────────────
- 14.07 ┤                             ╭──╯
- 13.13 ┤                         ╭───╯
- 12.20 ┤                      ╭──╯
- 11.27 ┤                   ╭──╯
- 10.33 ┤                ╭──╯
-  9.40 ┤            ╭───╯
-  8.47 ┤         ╭──╯
-  7.53 ┤         │
-  6.60 ┤         │
-  5.67 ┤         │
-  4.73 ┤         │
-  3.80 ┤      ╭──╯
-  2.87 ┤      │
-  1.93 ┤   ╭──╯
-  1.00 ┼───╯
-                                            replicas
+ 9.00 ┤                          ╭────────────────────────────────────────────────────
+ 8.47 ┤                          │
+ 7.93 ┤                      ╭───╯
+ 7.40 ┤                      │
+ 6.87 ┤                   ╭──╯
+ 6.33 ┤                   │
+ 5.80 ┤               ╭───╯
+ 5.27 ┤               │
+ 4.73 ┤            ╭──╯
+ 4.20 ┤         ╭──╯
+ 3.67 ┤         │
+ 3.13 ┤      ╭──╯
+ 2.60 ┤      │
+ 2.07 ┤   ╭──╯
+ 1.53 ┤   │
+ 1.00 ┼───╯
+                                           replicas
 initial store values: [s1=1] (stddev=0.00, mean=1.00, sum=1)
-last store values: [s1=15] (stddev=0.00, mean=15.00, sum=15)
+last store values: [s1=9] (stddev=0.00, mean=9.00, sum=9)
 
 # vim:ft=sh

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  //