Merge #144360

144360: cspann: convert search logic to use pull iteration r=mw5h a=andy-kimball

This PR fixes multiple issues that were causing the `TestIndexConcurrency`
test to fail. That test was no longer working after the change to create a new
non-transactional background split fixup.

Fixes: #143941

#### cspann: convert search logic to use pull iteration

Using non-transactional splits breaks certain assumptions that the C-SPANN
library relies upon:

* The K-means tree must be fully balanced - interior partitions can now be
  empty.
* Partition child keys are never duplicated - splitting can duplicates
  child partition keys, which can persist in the tree.
* Partition child keys always reference existing partitions - partition
  child keys can now reference missing partitions.
* Inserts will always trivially find an insertion partition - now, it's
  possible to get "blocked" trying to find a path to a target partition
  that supports inserts. The blockage could be in the form of an empty
  interior partition, a dangling child partition key, or a target
  partition that does not allow inserts.

To address these issues, this commit converts the existing search logic to
use "pull" iteration. Code that needs to search the tree can iteratively
get the next batch of results, and the next after that, and so on, without
knowing up front exactly how many results are needed. Each batch is sorted
by distance, with duplicates removed. However, batches are not strictly
ordered in relation to one another and duplicates can exist across batches
(though each subsequent batch does tend to have greater distances).

Pull iteration largely solves the issues noted above. For example, the
insert operation can pull one result at a time. If that partition does not
support inserts, it can pull the next. If a batch is empty due to hitting
a "dead end" in tree traversal, it can just pull the next batch.

Epic: CRDB-42943
Release note: None

#### cspann: update C-SPANN index code to tolerate missing partitions

Previously, when index operations like search or insert came across a
missing partition, they failed with an error, as this should never be
possible. However, with non-transactional splits, it's now acceptable
for partitions to be missing. This commit changes the index code to
tolerate this condition. It also updates transactional Store methods
to support the new semantics.

Epic: CRDB-42943
Release note: None

#### vecindex: remove Begin/Commit/Abort Transaction methods

Remove the BeginTransaction, CommitTransaction, and AbortTransaction methods
from the cspann.Store interface. These methods are superseded by the
RunTransaction method. Fix up all test usages of any of these methods to
use a new commontest.RunTransaction helper function.

Epic: CRDB-42943
Release note: None

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

Author: craig[bot]

pkg/cmd/vecbench/mem_provider.go

@@ -128,34 +128,27 @@ func (m *MemProvider) InsertVectors(
 func (m *MemProvider) Search(
 	ctx context.Context, vec vector.T, maxResults int, beamSize int, stats *cspann.SearchStats,
 ) (keys []cspann.KeyBytes, err error) {
-	var txn cspann.Txn
-	txn, err = m.store.BeginTransaction(ctx)
-	defer func() {
-		if err == nil {
-			err = m.store.CommitTransaction(ctx, txn)
-		}
+	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}
+		err = m.index.Search(ctx, &idxCtx, nil /* treeKey */, vec, &searchSet, searchOptions)
 		if err != nil {
-			err = m.store.AbortTransaction(ctx, txn)
+			return err
 		}
-	}()
-
-	// Search the store.
-	var idxCtx cspann.Context
-	idxCtx.Init(txn)
-	searchSet := cspann.SearchSet{MaxResults: maxResults}
-	searchOptions := cspann.SearchOptions{BaseBeamSize: beamSize}
-	err = m.index.Search(ctx, &idxCtx, nil /* treeKey */, vec, &searchSet, searchOptions)
-	if err != nil {
-		return nil, err
-	}
-	*stats = searchSet.Stats
+		*stats = searchSet.Stats
 
-	// Get result keys.
-	results := searchSet.PopResults()
-	keys = make([]cspann.KeyBytes, len(results))
-	for i, res := range results {
-		keys[i] = []byte(res.ChildKey.KeyBytes)
-	}
+		// Get result keys.
+		results := searchSet.PopResults()
+		keys = make([]cspann.KeyBytes, len(results))
+		for i, res := range results {
+			keys[i] = []byte(res.ChildKey.KeyBytes)
+		}
+
+		return nil
+	})
 
 	return keys, err
 }

pkg/sql/vecindex/cspann/BUILD.bazel

@@ -23,6 +23,7 @@ go_library(
         "partition.go",
         "partition_metadata.go",
         "search_set.go",
+        "searcher.go",
         "split_data.go",
         "store.go",
         "store_errors.go",

pkg/sql/vecindex/cspann/childkey_dedup.go

@@ -15,10 +15,10 @@ import (
 // hashKeyFunc is a function type for hashing KeyBytes.
 type hashKeyFunc func(KeyBytes) uint64
 
-// childKeyDeDup provides de-duplication for ChildKey values. It supports both
+// ChildKeyDeDup provides de-duplication for ChildKey values. It supports both
 // PartitionKey and KeyBytes child keys efficiently without making unnecessary
 // allocations.
-type childKeyDeDup struct {
+type ChildKeyDeDup struct {
 	// initialCapacity is used to initialize the size of the data structures used
 	// by the de-duplicator.
 	initialCapacity int
@@ -39,17 +39,22 @@ type childKeyDeDup struct {
 }
 
 // Init initializes the de-duplicator.
-func (dd *childKeyDeDup) Init(capacity int) {
+func (dd *ChildKeyDeDup) Init(capacity int) {
 	dd.initialCapacity = capacity
 	dd.seed = maphash.MakeSeed()
 	dd.hashKeyBytes = dd.defaultHashKeyBytes
 	dd.Clear()
 }
 
+// Count returns the number of keys in the de-duplicator.
+func (dd *ChildKeyDeDup) Count() int {
+	return len(dd.partitionKeys) + len(dd.keyBytesMap)
+}
+
 // TryAdd attempts to add a child key to the deduplication set. It returns true
 // if the key was added (wasn't a duplicate), or false if the key already exists
 // (is a duplicate).
-func (dd *childKeyDeDup) TryAdd(childKey ChildKey) bool {
+func (dd *ChildKeyDeDup) TryAdd(childKey ChildKey) bool {
 	// Handle PartitionKey case - simple map lookup.
 	if childKey.PartitionKey != 0 {
 		// Lazily initialize the partitionKeys map.
@@ -102,19 +107,19 @@ func (dd *childKeyDeDup) TryAdd(childKey ChildKey) bool {
 }
 
 // Clear removes all entries from the deduplication set.
-func (dd *childKeyDeDup) Clear() {
+func (dd *ChildKeyDeDup) Clear() {
 	// Reset all the data structures.
 	clear(dd.partitionKeys)
 	clear(dd.keyBytesMap)
 }
 
 // defaultHashKeyBytes is the default implementation of hashKeyBytes.
-func (dd *childKeyDeDup) defaultHashKeyBytes(key KeyBytes) uint64 {
+func (dd *ChildKeyDeDup) defaultHashKeyBytes(key KeyBytes) uint64 {
 	return maphash.Bytes(dd.seed, key)
 }
 
 // rehash creates a new hash from an existing hash to resolve collisions.
-func (dd *childKeyDeDup) rehash(hash uint64) uint64 {
+func (dd *ChildKeyDeDup) rehash(hash uint64) uint64 {
 	// These constants are large 64-bit primes.
 	hash ^= 0xc3a5c85c97cb3127
 	hash ^= hash >> 33

pkg/sql/vecindex/cspann/childkey_dedup_test.go

@@ -19,42 +19,48 @@ func TestChildKeyDeDupAddPartitionKey(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	var dd childKeyDeDup
+	var dd ChildKeyDeDup
 	dd.Init(10)
 
 	// Add a new PartitionKey.
 	added := dd.TryAdd(ChildKey{PartitionKey: 123})
 	require.True(t, added)
+	require.Equal(t, 1, dd.Count())
 
 	// Try to add the same key again (should be a duplicate).
 	added = dd.TryAdd(ChildKey{PartitionKey: 123})
 	require.False(t, added)
+	require.Equal(t, 1, dd.Count())
 
 	// Add a different PartitionKey.
 	added = dd.TryAdd(ChildKey{PartitionKey: 456})
 	require.True(t, added)
+	require.Equal(t, 2, dd.Count())
 }
 
 func TestChildKeyDeDupAddKeyBytes(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	var dd childKeyDeDup
+	var dd ChildKeyDeDup
 	dd.Init(10)
 
 	// Add a new KeyBytes.
 	key1 := []byte("key1")
 	added := dd.TryAdd(ChildKey{KeyBytes: key1})
 	require.True(t, added)
+	require.Equal(t, 1, dd.Count())
 
 	// Try to add the same key again (should be a duplicate).
 	added = dd.TryAdd(ChildKey{KeyBytes: key1})
 	require.False(t, added)
+	require.Equal(t, 1, dd.Count())
 
 	// Add a different KeyBytes.
 	key2 := []byte("key2")
 	added = dd.TryAdd(ChildKey{KeyBytes: key2})
 	require.True(t, added)
+	require.Equal(t, 2, dd.Count())
 
 	// Verify that both keys are properly stored by checking for duplicates.
 	added = dd.TryAdd(ChildKey{KeyBytes: key1})
@@ -67,17 +73,19 @@ func TestChildKeyDeDupClear(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	var dd childKeyDeDup
+	var dd ChildKeyDeDup
 	dd.Init(10)
 
 	// Add a mix of keys.
 	require.True(t, dd.TryAdd(ChildKey{PartitionKey: 123}))
 	require.True(t, dd.TryAdd(ChildKey{KeyBytes: []byte("key1")}))
+	require.Equal(t, 2, dd.Count())
 
 	// Clear the deduplicator.
 	dd.Clear()
 
 	// Verify keys were cleared.
+	require.Equal(t, 0, dd.Count())
 	require.True(t, dd.TryAdd(ChildKey{PartitionKey: 123}))
 	require.True(t, dd.TryAdd(ChildKey{KeyBytes: []byte("key1")}))
 }
@@ -96,7 +104,7 @@ func TestChildKeyDeDupRehashing(t *testing.T) {
 		return maphash.Bytes(seed, key)
 	}
 
-	var dd childKeyDeDup
+	var dd ChildKeyDeDup
 	dd.Init(100)
 	dd.hashKeyBytes = customHashFunc