Merge #144525

144525: vecindex: add settings and storage options for vector indexes r=drewkimball,mw5h a=andy-kimball

#### vecindex: disable vector indexes by default

For 25.2, vector indexes are shipping as a public preview, but disabled until
users set the cluster setting feature.vector_index.enabled=true. Creating a
vector index gives an error until then.

Epic: CRDB-42943
Release note: None

#### vecindex: add vector-search-beam-size session setting

Add a new "vector-search-beam-size" session setting that controls the
breadth of the vector index search. The higher the value, the more
accurate is the search, but the more processing it requires.

Epic: CRDB-42943
Release note: None

#### vecindex: offer vector index storage params

Add support for several storage parameters when creating a vector index:

* build_beam_size: controls accuracy of index build
* min_partition_size: min size of partition without being merged
* max_partition_size: max size of partition without being split

Here is an example:

CREATE VECTOR INDEX ON storage_params (v)
WITH (min_partition_size = 8, max_partition_size = 64)

Epic: CRDB-42943
Release note: None

#### vecindex: hook up StalledOpTimeout

Hook up the vector index's stalled op timeout callback to the cluster
setting. While the setting was added in a previous PR, it was never
being used.

Epic: CRDB-42943
Release note: None

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

Author: craig[bot]

docs/generated/settings/settings-for-tenants.txt

@@ -75,6 +75,7 @@ feature.import.enabled	boolean	true	set to true to enable imports, false to disa
 feature.restore.enabled	boolean	true	set to true to enable restore, false to disable; default is true	application
 feature.schema_change.enabled	boolean	true	set to true to enable schema changes, false to disable; default is true	application
 feature.stats.enabled	boolean	true	set to true to enable CREATE STATISTICS/ANALYZE, false to disable; default is true	application
+feature.vector_index.enabled	boolean	false	set to true to enable vector indexes, false to disable; default is false	application
 jobs.retention_time	duration	336h0m0s	the amount of time for which records for completed jobs are retained	application
 kv.bulk_sst.target_size	byte size	16 MiB	target size for SSTs emitted from export requests; export requests (i.e. BACKUP) may buffer up to the sum of kv.bulk_sst.target_size and kv.bulk_sst.max_allowed_overage in memory	system-visible
 kv.closed_timestamp.follower_reads.enabled (alias: kv.closed_timestamp.follower_reads_enabled)	boolean	true	allow (all) replicas to serve consistent historical reads based on closed timestamp information	system-visible

docs/generated/settings/settings.html

@@ -80,6 +80,7 @@
 <tr><td><div id="setting-feature-restore-enabled" class="anchored"><code>feature.restore.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>set to true to enable restore, false to disable; default is true</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-feature-schema-change-enabled" class="anchored"><code>feature.schema_change.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>set to true to enable schema changes, false to disable; default is true</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-feature-stats-enabled" class="anchored"><code>feature.stats.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>set to true to enable CREATE STATISTICS/ANALYZE, false to disable; default is true</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
+<tr><td><div id="setting-feature-vector-index-enabled" class="anchored"><code>feature.vector_index.enabled</code></div></td><td>boolean</td><td><code>false</code></td><td>set to true to enable vector indexes, false to disable; default is false</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-jobs-retention-time" class="anchored"><code>jobs.retention_time</code></div></td><td>duration</td><td><code>336h0m0s</code></td><td>the amount of time for which records for completed jobs are retained</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-lease-rebalance-threshold" class="anchored"><code>kv.allocator.lease_rebalance_threshold</code></div></td><td>float</td><td><code>0.05</code></td><td>minimum fraction away from the mean a store&#39;s lease count can be before it is considered for lease-transfers</td><td>Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-load-based-lease-rebalancing-enabled" class="anchored"><code>kv.allocator.load_based_lease_rebalancing.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>set to enable rebalancing of range leases based on load and latency</td><td>Dedicated/Self-Hosted</td></tr>

pkg/cmd/vecbench/main.go

@@ -65,9 +65,9 @@ var flagHideProgress = flag.Bool("hide-progress", false, "Hide progress during i
 // Search command options.
 var flagMaxResults = flag.Int("k", 10, "Number of search results, used in recall calculation.")
 var flagBeamSize = flag.Int(
-	"default-beam-size",
+	"build-beam-size",
 	8,
-	"Default beam size used for building and searching the index.")
+	"Default beam size used for building the index.")
 var flagSearchBeamSizes = flag.String(
 	"search-beam-sizes",
 	"1,2,4,8,16,32,64,128,256,512",
@@ -248,30 +248,37 @@ func (vb *vectorBench) SearchIndex() {
 	}
 
 	doSearch := func(beamSize int) {
+		// Prepare to search.
+		maxResults := *flagMaxResults
+		state, err := vb.provider.SetupSearch(vb.ctx, maxResults, beamSize)
+		if err != nil {
+			panic(err)
+		}
+
 		start := timeutil.Now()
 
 		// Search for test vectors.
-		var sumMAP, sumVectors, sumLeafVectors, sumFullVectors, sumPartitions float64
+		var sumRecall, sumVectors, sumLeafVectors, sumFullVectors, sumPartitions float64
 		count := data.Test.Count
-		for i := 0; i < count; i++ {
+		for i := range count {
 			// Calculate truth set for the vector.
 			queryVector := data.Test.At(i)
 
 			var stats cspann.SearchStats
-			prediction, err := vb.provider.Search(vb.ctx, queryVector, *flagMaxResults, beamSize, &stats)
+			prediction, err := vb.provider.Search(vb.ctx, state, queryVector, &stats)
 			if err != nil {
 				panic(err)
 			}
 
-			primaryKeys := make([]byte, len(prediction)*4)
-			truth := make([]cspann.KeyBytes, len(prediction))
-			for neighbor := 0; neighbor < len(prediction); neighbor++ {
+			primaryKeys := make([]byte, maxResults*4)
+			truth := make([]cspann.KeyBytes, maxResults)
+			for neighbor := range maxResults {
 				primaryKey := primaryKeys[neighbor*4 : neighbor*4+4]
 				binary.BigEndian.PutUint32(primaryKey, uint32(data.Neighbors[i][neighbor]))
 				truth[neighbor] = primaryKey
 			}
 
-			sumMAP += findMAP(prediction, truth)
+			sumRecall += calculateRecall(prediction, truth)
 			sumVectors += float64(stats.QuantizedVectorCount)
 			sumLeafVectors += float64(stats.QuantizedLeafVectorCount)
 			sumFullVectors += float64(stats.FullVectorCount)
@@ -280,7 +287,7 @@ func (vb *vectorBench) SearchIndex() {
 
 		elapsed := timeutil.Since(start)
 		fmt.Printf("%d\t%0.2f%%\t%0.0f\t%0.0f\t%0.2f\t%0.2f\t%0.2f\n",
-			beamSize, sumMAP/float64(count)*100,
+			beamSize, sumRecall/float64(count)*100,
 			sumLeafVectors/float64(count), sumVectors/float64(count),
 			sumFullVectors/float64(count), sumPartitions/float64(count),
 			float64(count)/elapsed.Seconds())
@@ -289,14 +296,14 @@ func (vb *vectorBench) SearchIndex() {
 	fmt.Println()
 	fmt.Printf(White+"%s\n"+Reset, vb.datasetName)
 	fmt.Printf(
-		White+"%d train vectors, %d test vectors, %d dimensions, %d/%d min/max partitions, base beam size %d\n"+Reset,
+		White+"%d train vectors, %d test vectors, %d dimensions, %d/%d min/max partitions, build beam size %d\n"+Reset,
 		data.Count, data.Test.Count, data.Test.Dims,
 		minPartitionSize, maxPartitionSize, *flagBeamSize)
 	fmt.Println(vb.provider.FormatStats())
 
 	fmt.Printf("beam\trecall\tleaf\tall\tfull\tpartns\tqps\n")
 
-	// Search multiple times with different beam sizes.
+	// Search multiple times with different search beam sizes.
 	beamSizeStrs := strings.Split(*flagSearchBeamSizes, ",")
 	for i := range beamSizeStrs {
 		beamSize, err := strconv.Atoi(beamSizeStrs[i])
@@ -626,15 +633,11 @@ func loadDataset(fileName string) dataset {
 	return data
 }
 
-// findMAP returns mean average precision, which compares a set of predicted
-// results with the true set of results. Both sets are expected to be of equal
-// length. It returns the percentage overlap of the predicted set with the truth
-// set.
-func findMAP(prediction, truth []cspann.KeyBytes) float64 {
-	if len(prediction) != len(truth) {
-		panic(errors.AssertionFailedf("prediction and truth sets are not same length"))
-	}
-
+// calculateRecall returns the percentage overlap of the predicted set with the
+// truth set. If the predicted set has fewer items than the truth set, it is
+// treated as if the predicted set has missing/incorrect items that reduce the
+// recall rate.
+func calculateRecall(prediction, truth []cspann.KeyBytes) float64 {
 	predictionMap := make(map[string]bool, len(prediction))
 	for _, p := range prediction {
 		predictionMap[string(p)] = true

pkg/cmd/vecbench/mem_provider.go

@@ -23,6 +23,16 @@ import (
 
 const seed = 42
 
+// MemSearchState holds prepared search state.
+type MemSearchState struct {
+	beamSize   int
+	maxResults int
+}
+
+// Close is a no-op for the MemProvider.
+func (s *MemSearchState) Close() {
+}
+
 // MemProvider implements VectorProvider using an in-memory store.
 type MemProvider struct {
 	stopper          *stop.Stopper
@@ -124,16 +134,24 @@ func (m *MemProvider) InsertVectors(
 	})
 }
 
+// SetupSearch implements the VectorProvider interface.
+func (m *MemProvider) SetupSearch(
+	ctx context.Context, maxResults int, beamSize int,
+) (SearchState, error) {
+	return &MemSearchState{maxResults: maxResults, beamSize: beamSize}, nil
+}
+
 // Search implements the VectorProvider interface.
 func (m *MemProvider) Search(
-	ctx context.Context, vec vector.T, maxResults int, beamSize int, stats *cspann.SearchStats,
+	ctx context.Context, state SearchState, vec vector.T, stats *cspann.SearchStats,
 ) (keys []cspann.KeyBytes, err error) {
+	memState := state.(*MemSearchState)
 	err = m.store.RunTransaction(ctx, func(txn cspann.Txn) error {
 		// Search the store.
 		var idxCtx cspann.Context
 		idxCtx.Init(txn)
-		searchSet := cspann.SearchSet{MaxResults: maxResults}
-		searchOptions := cspann.SearchOptions{BaseBeamSize: beamSize}
+		searchSet := cspann.SearchSet{MaxResults: memState.maxResults}
+		searchOptions := cspann.SearchOptions{BaseBeamSize: memState.beamSize}
 		err = m.index.Search(ctx, &idxCtx, nil /* treeKey */, vec, &searchSet, searchOptions)
 		if err != nil {
 			return err

pkg/cmd/vecbench/sql_provider.go

@@ -24,6 +24,21 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 )
 
+// SQLSearchState holds an open connection and a prepared statement to run on
+// that connection in repeated calls to Search.
+type SQLSearchState struct {
+	conn  *pgxpool.Conn
+	query string
+}
+
+// Close releases the connection back to the pool.
+func (s *SQLSearchState) Close() {
+	if s.conn != nil {
+		s.conn.Release()
+		s.conn = nil
+	}
+}
+
 // SQLProvider implements VectorProvider using a SQL database connection to a
 // CockroachDB instance.
 type SQLProvider struct {
@@ -135,7 +150,7 @@ func (s *SQLProvider) InsertVectors(
 			batch := &pgx.Batch{}
 
 			// Insert vectors in batches.
-			for i := 0; i < vectors.Count; i++ {
+			for i := range vectors.Count {
 				batch.Queue(fmt.Sprintf(
 					"INSERT INTO %s (id, embedding) VALUES ($1, $2)",
 					s.tableName),
@@ -165,20 +180,54 @@ func (s *SQLProvider) InsertVectors(
 	}
 }
 
-// Search implements the VectorProvider interface.
-func (s *SQLProvider) Search(
-	ctx context.Context, vec vector.T, maxResults int, beamSize int, stats *cspann.SearchStats,
-) ([]cspann.KeyBytes, error) {
-	// Construct a query that searches for similar vectors.
+// SetupSearch implements the VectorProvider interface.
+func (s *SQLProvider) SetupSearch(
+	ctx context.Context, maxResults int, beamSize int,
+) (SearchState, error) {
+	// Acquire a connection from the pool.
+	conn, err := s.pool.Acquire(ctx)
+	if err != nil {
+		return nil, errors.Wrap(err, "acquiring connection from pool")
+	}
+	defer func() {
+		if conn != nil {
+			conn.Release()
+		}
+	}()
+
+	// Set the vector_search_beam_size session variable.
+	_, err = conn.Exec(ctx, fmt.Sprintf("SET vector_search_beam_size = %d", beamSize))
+	if err != nil {
+		return nil, errors.Wrap(err, "setting vector_search_beam_size")
+	}
+
+	// Construct the query for vector search.
 	query := fmt.Sprintf(`
 		SELECT id
 		FROM %s
 		ORDER BY embedding <-> $1
-		LIMIT $2
-	`, s.tableName)
+		LIMIT %d
+	`, s.tableName, maxResults)
+
+	state := &SQLSearchState{
+		conn:  conn,
+		query: query,
+	}
+	conn = nil
+	return state, nil
+}
+
+// Search implements the VectorProvider interface.
+func (s *SQLProvider) Search(
+	ctx context.Context, state SearchState, vec vector.T, stats *cspann.SearchStats,
+) ([]cspann.KeyBytes, error) {
+	sqlState, ok := state.(*SQLSearchState)
+	if !ok {
+		return nil, errors.New("invalid search state type")
+	}
 
-	// Execute the query.
-	rows, err := s.pool.Query(ctx, query, vec, maxResults)
+	// Execute the prepared statement.
+	rows, err := sqlState.conn.Query(ctx, sqlState.query, vec)
 	if err != nil {
 		return nil, errors.Wrap(err, "executing search query")
 	}

pkg/cmd/vecbench/vector_provider.go

@@ -21,6 +21,13 @@ type IndexMetric struct {
 	Value float64
 }
 
+// SearchState holds prepared state needed by a vector provider across multiple
+// calls to Search.
+type SearchState interface {
+	// Close releases any resources held by the search state.
+	Close()
+}
+
 // VectorProvider abstracts the operations needed for vector storage and
 // retrieval. This allows different implementations (in-memory, SQL-based, etc.)
 // to provide the functionality needed by vecbench.
@@ -47,11 +54,16 @@ type VectorProvider interface {
 	// identified by a key.
 	InsertVectors(ctx context.Context, keys []cspann.KeyBytes, vectors vector.Set) error
 
+	// SetupSearch allows the provider to perform expensive up-front steps in
+	// preparation for many calls to Search. It returns provider-specific state
+	// that will be passed to Search.
+	SetupSearch(ctx context.Context, maxResults int, beamSize int) (SearchState, error)
+
 	// Search searches for vectors similar to the query vector. It returns the
 	// keys of the most similar vectors. If supported, stats are recorded in
 	// "stats" for this search.
 	Search(
-		ctx context.Context, vec vector.T, maxResults int, beamSize int, stats *cspann.SearchStats,
+		ctx context.Context, state SearchState, vec vector.T, stats *cspann.SearchStats,
 	) ([]cspann.KeyBytes, error)
 
 	// GetMetrics returns interesting metrics for the vector index. Each provider

pkg/sql/backfill/backfill_test.go

@@ -51,6 +51,9 @@ func TestVectorColumnAndIndexBackfill(t *testing.T) {
 	defer srv.Stopper().Stop(ctx)
 	sqlDB := sqlutils.MakeSQLRunner(db)
 
+	// Enable vector indexes.
+	sqlDB.Exec(t, `SET CLUSTER SETTING feature.vector_index.enabled = true`)
+
 	// Create a table with a vector column
 	sqlDB.Exec(t, `
 		CREATE TABLE vectors (
@@ -115,6 +118,9 @@ func TestConcurrentOperationsDuringVectorIndexCreation(t *testing.T) {
 	defer srv.Stopper().Stop(ctx)
 	sqlDB := sqlutils.MakeSQLRunner(db)
 
+	// Enable vector indexes.
+	sqlDB.Exec(t, `SET CLUSTER SETTING feature.vector_index.enabled = true`)
+
 	// Create a table with a vector column
 	sqlDB.Exec(t, `
 		CREATE TABLE vectors (

pkg/sql/catalog/catformat/index.go

@@ -275,6 +275,18 @@ func formatStorageConfigs(
 	table catalog.TableDescriptor, index *descpb.IndexDescriptor, f *tree.FmtCtx,
 ) error {
 	numCustomSettings := 0
+	writeCustomSetting := func(key, val string) {
+		if numCustomSettings > 0 {
+			f.WriteString(", ")
+		} else {
+			f.WriteString(" WITH (")
+		}
+		numCustomSettings++
+		f.WriteString(key)
+		f.WriteString("=")
+		f.WriteString(val)
+	}
+
 	if index.GeoConfig.S2Geometry != nil || index.GeoConfig.S2Geography != nil {
 		var s2Config *geopb.S2Config
 
@@ -297,15 +309,7 @@ func formatStorageConfigs(
 				{`s2_max_cells`, s2Config.MaxCells, defaultS2Config.MaxCells},
 			} {
 				if check.val != check.defaultVal {
-					if numCustomSettings > 0 {
-						f.WriteString(", ")
-					} else {
-						f.WriteString(" WITH (")
-					}
-					numCustomSettings++
-					f.WriteString(check.key)
-					f.WriteString("=")
-					f.WriteString(strconv.Itoa(int(check.val)))
+					writeCustomSetting(check.key, strconv.Itoa(int(check.val)))
 				}
 			}
 		}
@@ -332,29 +336,26 @@ func formatStorageConfigs(
 				{`geometry_max_y`, cfg.MaxY, defaultConfig.S2Geometry.MaxY},
 			} {
 				if check.val != check.defaultVal {
-					if numCustomSettings > 0 {
-						f.WriteString(", ")
-					} else {
-						f.WriteString(" WITH (")
-					}
-					numCustomSettings++
-					f.WriteString(check.key)
-					f.WriteString("=")
-					f.WriteString(strconv.FormatFloat(check.val, 'f', -1, 64))
+					writeCustomSetting(check.key, strconv.FormatFloat(check.val, 'f', -1, 64))
 				}
 			}
 		}
 	}
 
-	if index.IsSharded() {
-		if numCustomSettings > 0 {
-			f.WriteString(", ")
-		} else {
-			f.WriteString(" WITH (")
+	if index.Type == idxtype.VECTOR {
+		if index.VecConfig.BuildBeamSize != 0 {
+			writeCustomSetting(`build_beam_size`, strconv.Itoa(int(index.VecConfig.BuildBeamSize)))
 		}
-		f.WriteString(`bucket_count=`)
-		f.WriteString(strconv.FormatInt(int64(index.Sharded.ShardBuckets), 10))
-		numCustomSettings++
+		if index.VecConfig.MinPartitionSize != 0 {
+			writeCustomSetting(`min_partition_size`, strconv.Itoa(int(index.VecConfig.MinPartitionSize)))
+		}
+		if index.VecConfig.MaxPartitionSize != 0 {
+			writeCustomSetting(`max_partition_size`, strconv.Itoa(int(index.VecConfig.MaxPartitionSize)))
+		}
+	}
+
+	if index.IsSharded() {
+		writeCustomSetting(`bucket_count`, strconv.FormatInt(int64(index.Sharded.ShardBuckets), 10))
 	}
 
 	if numCustomSettings > 0 {