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

Author: andy-kimball

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
 

pkg/sql/vecindex/cspann/commontest/storetests.go

@@ -1085,18 +1085,17 @@ func (suite *StoreTestSuite) TestTryRemoveFromPartition() {
 		suite.NoError(err)
 		suite.False(removed) // no vectors removed
 
-		// Try to remove the last remaining vectors in the partition.
+		// Remove the last remaining vectors in the partition.
 		childKeys = []cspann.ChildKey{partitionKey2, partitionKey3}
 		removed, err = store.TryRemoveFromPartition(suite.ctx, treeKey, partitionKey,
 			childKeys, expected)
-		suite.ErrorContains(err, "cannot remove last remaining vector from non-leaf partition")
+		suite.NoError(err)
 		suite.True(removed)
 
-		// Verify that partition contains one last remaining vector.
+		// Verify that partition contains no remaining vectors.
 		partition, err = store.TryGetPartition(suite.ctx, treeKey, partitionKey)
 		suite.NoError(err)
-		suite.Equal(1, partition.Count())
-		suite.Equal([]cspann.ChildKey{partitionKey3}, partition.ChildKeys())
+		suite.Equal(0, partition.Count())
 	}
 
 	suite.Run("default tree", func() {
@@ -1347,7 +1346,7 @@ func (suite *StoreTestSuite) testLeafPartition(
 			// Search partition.
 			searchSet := cspann.SearchSet{MaxResults: 2}
 			toSearch := []cspann.PartitionToSearch{{Key: partitionKey}}
-			 err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{1, 1}, &searchSet)
+			err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{1, 1}, &searchSet)
 			suite.NoError(err)
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 1, ErrorBound: 0,

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -124,6 +124,7 @@ func (fw *fixupWorker) Start(ctx context.Context) {
 		}
 		fw.treeKey = next.TreeKey
 		fw.singleStep = next.SingleStep
+		fw.txn = nil
 
 		// Invoke the fixup function. Note that we do not hold the lock while
 		// processing the fixup.
@@ -574,15 +575,10 @@ func (fw *fixupWorker) linkNearbyVectors(
 	// TODO(andyk): Add way to filter search set in order to skip vectors deeper
 	// down in the search rather than afterwards.
 	idxCtx := fw.reuseIndexContext(fw.txn, fw.treeKey)
-	idxCtx.options = SearchOptions{ReturnVectors: true}
+	idxCtx.options = SearchOptions{}
 	idxCtx.level = partition.Level()
 	idxCtx.randomized = partition.Centroid()
 
-	// Ensure that the search never returns the last remaining vector in a
-	// non-leaf partition, in order to avoid moving it and creating an empty
-	// non-leaf partition, which is not allowed by a balanced K-means tree.
-	idxCtx.ignoreLonelyVector = partition.Level() != LeafLevel
-
 	// Don't link more vectors than the number of remaining slots in the split
 	// partition, to avoid triggering another split.
 	maxResults := fw.index.options.MaxPartitionSize - partition.Count()
@@ -828,10 +824,13 @@ func (fw *fixupWorker) getFullVectorsForPartition(
 		return nil
 	}
 
- 	// Run in a transaction if not already.
+	// Run in a transaction if not already.
 	if fw.txn == nil {
 		err = fw.index.store.RunTransaction(ctx, func(txn Txn) error {
 			fw.txn = txn
+			defer func() {
+				fw.txn = nil
+			}()
 			return run()
 		})
 	} else {