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

Author: andy-kimball

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

@@ -111,7 +111,7 @@ func (suite *StoreTestSuite) TestRunTransaction() {
 	suite.ErrorContains(store.RunTransaction(suite.ctx, func(tx cspann.Txn) error {
 		toSearch := []cspann.PartitionToSearch{{Key: cspann.RootKey}}
 		searchSet := cspann.SearchSet{MaxResults: 1}
-		_, err := tx.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{1, -1}, &searchSet)
+		err := tx.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{1, -1}, &searchSet)
 		suite.NoError(err)
 		suite.Equal(1, toSearch[0].Count)
 		return errors.New("abort")
@@ -260,16 +260,15 @@ func (suite *StoreTestSuite) TestNonRootPartition() {
 			// Search partition.
 			toSearch := []cspann.PartitionToSearch{{Key: partitionKey}}
 			searchSet := cspann.SearchSet{MaxResults: 1}
-			searchLevel, err := txn.SearchPartitions(
-				suite.ctx, treeKey, toSearch, vector.T{5, -1}, &searchSet)
+			err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{5, -1}, &searchSet)
 			suite.NoError(err)
-			suite.Equal(cspann.SecondLevel, searchLevel)
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 17, ErrorBound: 0, CentroidDistance: 0,
 				ParentPartitionKey: partitionKey, ChildKey: partitionKey3, ValueBytes: valueBytes3}
 			results := searchSet.PopResults()
 			RoundResults(results, 4)
 			suite.Equal(cspann.SearchResults{result1}, results)
+			suite.Equal(cspann.SecondLevel, toSearch[0].Level)
 			suite.Equal(2, toSearch[0].Count)
 		})
 	}
@@ -308,17 +307,18 @@ func (suite *StoreTestSuite) TestSearchMultiplePartitions() {
 
 			searchSet := cspann.SearchSet{MaxResults: 2}
 			toSearch := []cspann.PartitionToSearch{{Key: cspann.RootKey}, {Key: partitionKey}}
-			level, err := txn.SearchPartitions(suite.ctx, treeKey, toSearch, vec4, &searchSet)
+			err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vec4, &searchSet)
 			suite.NoError(err)
-			suite.Equal(cspann.LeafLevel, level)
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 24, ErrorBound: 24.08, CentroidDistance: 3.16,
 				ParentPartitionKey: partitionKey, ChildKey: primaryKey1, ValueBytes: valueBytes1}
 			result2 := cspann.SearchResult{
 				QuerySquaredDistance: 29, ErrorBound: 0, CentroidDistance: 5,
 				ParentPartitionKey: cspann.RootKey, ChildKey: primaryKey3, ValueBytes: valueBytes3}
 			suite.Equal(cspann.SearchResults{result1, result2}, RoundResults(searchSet.PopResults(), 2))
+			suite.Equal(cspann.LeafLevel, toSearch[0].Level)
 			suite.Equal(3, toSearch[0].Count)
+			suite.Equal(cspann.LeafLevel, toSearch[1].Level)
 			suite.Equal(2, toSearch[1].Count)
 		})
 	}
@@ -602,6 +602,7 @@ func (suite *StoreTestSuite) TestGetFullVectors() {
 				{Key: cspann.ChildKey{KeyBytes: cspann.KeyBytes{0}}},
 				{Key: cspann.ChildKey{PartitionKey: partitionKey}},
 				{Key: cspann.ChildKey{KeyBytes: key3}},
+				{Key: cspann.ChildKey{PartitionKey: cspann.PartitionKey(99)}}, // No such partition.
 			}
 			err := txn.GetFullVectors(suite.ctx, treeKey, results)
 			suite.NoError(err)
@@ -612,6 +613,7 @@ func (suite *StoreTestSuite) TestGetFullVectors() {
 			suite.Nil(results[4].Vector)
 			suite.Equal(vector.T{4, 3}, results[5].Vector)
 			suite.Equal(vec3, results[6].Vector)
+			suite.Nil(results[7].Vector)
 
 			// Grab another set of vectors to ensure that saved state is properly reset.
 			results = []cspann.VectorWithKey{
@@ -1228,10 +1230,10 @@ func (suite *StoreTestSuite) testEmptyOrMissingRoot(store TestStore, treeID int,
 		RunTransaction(suite.ctx, suite.T(), store, func(txn cspann.Txn) {
 			searchSet := cspann.SearchSet{MaxResults: 2}
 			toSearch := []cspann.PartitionToSearch{{Key: cspann.RootKey}}
-			level, 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)
-			suite.Equal(cspann.LeafLevel, level)
 			suite.Nil(searchSet.PopResults())
+			suite.Equal(cspann.LeafLevel, toSearch[0].Level)
 			suite.Equal(0, toSearch[0].Count)
 		})
 	})
@@ -1345,10 +1347,8 @@ func (suite *StoreTestSuite) testLeafPartition(
 			// Search partition.
 			searchSet := cspann.SearchSet{MaxResults: 2}
 			toSearch := []cspann.PartitionToSearch{{Key: partitionKey}}
-			searchLevel, 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)
-			suite.Equal(cspann.LeafLevel, searchLevel)
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 1, ErrorBound: 0,
 				CentroidDistance:   testutils.RoundFloat(num32.L2Distance(vec1, centroid), 4),
@@ -1367,6 +1367,7 @@ func (suite *StoreTestSuite) testLeafPartition(
 			results := searchSet.PopResults()
 			RoundResults(results, 4)
 			suite.Equal(cspann.SearchResults{result1, result2}, results)
+			suite.Equal(cspann.LeafLevel, toSearch[0].Level)
 			suite.Equal(3, toSearch[0].Count)
 
 			// Ensure partition metadata is updated.
@@ -1407,10 +1408,8 @@ func (suite *StoreTestSuite) testLeafPartition(
 			// Search partition.
 			searchSet := cspann.SearchSet{MaxResults: 1}
 			toSearch := []cspann.PartitionToSearch{{Key: partitionKey}}
-			searchLevel, err := txn.SearchPartitions(
-				suite.ctx, treeKey, toSearch, vector.T{10, -5}, &searchSet)
+			err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{10, -5}, &searchSet)
 			suite.NoError(err)
-			suite.Equal(cspann.LeafLevel, searchLevel)
 			result1 := cspann.SearchResult{
 				QuerySquaredDistance: 25, ErrorBound: 0,
 				CentroidDistance:   testutils.RoundFloat(num32.L2Distance(vec4, centroid), 4),
@@ -1423,6 +1422,7 @@ func (suite *StoreTestSuite) testLeafPartition(
 			results := searchSet.PopResults()
 			RoundResults(results, 4)
 			suite.Equal(cspann.SearchResults{result1}, results)
+			suite.Equal(cspann.LeafLevel, toSearch[0].Level)
 			suite.Equal(3, toSearch[0].Count)
 		})
 	})
@@ -1470,16 +1470,15 @@ func (suite *StoreTestSuite) setRootPartition(store TestStore, treeID int) {
 		// Search partition.
 		searchSet := cspann.SearchSet{MaxResults: 1}
 		toSearch := []cspann.PartitionToSearch{{Key: cspann.RootKey}}
-		searchLevel, err := txn.SearchPartitions(
-			suite.ctx, treeKey, toSearch, vector.T{5, 5}, &searchSet)
+		err = txn.SearchPartitions(suite.ctx, treeKey, toSearch, vector.T{5, 5}, &searchSet)
 		suite.NoError(err)
-		suite.Equal(cspann.SecondLevel, searchLevel)
 		result1 := cspann.SearchResult{
 			QuerySquaredDistance: 5, ErrorBound: 0, CentroidDistance: 3.1623,
 			ParentPartitionKey: cspann.RootKey, ChildKey: partitionKey2, ValueBytes: valueBytes2}
 		results := searchSet.PopResults()
 		RoundResults(results, 4)
 		suite.Equal(cspann.SearchResults{result1}, results)
+		suite.Equal(cspann.SecondLevel, toSearch[0].Level)
 		suite.Equal(2, toSearch[0].Count)
 	})
 }

pkg/sql/vecindex/cspann/fixup_worker.go

@@ -758,6 +758,10 @@ func (fw *fixupWorker) deleteVector(
 				// Root partition's level has been updated, so just abort.
 				return nil
 			} else if err != nil {
+				if errors.Is(err, ErrPartitionNotFound) {
+					log.VEventf(ctx, 2, "partition %d no longer exists, do not delete vector", partitionKey)
+					return nil
+				}
 				return errors.Wrapf(err, "getting root partition's level")
 			}
 		}

pkg/sql/vecindex/cspann/index.go

@@ -605,6 +605,21 @@ func (vi *Index) searchForInsertHelper(
 	var seen map[PartitionKey]struct{}
 	maxResults := 2
 	allowRetry := true
+
+	// Check the next partition in the result list. If there are no more
+	// partitions to check, then expand the search to a wider set of partitions
+	// by increasing MaxResults.
+	checkNextPartition := func(partitionKey PartitionKey) {
+		results = results[1:]
+		if len(results) == 0 {
+			if seen == nil {
+				seen = make(map[PartitionKey]struct{})
+			}
+			seen[partitionKey] = struct{}{}
+			maxResults *= 2
+		}
+	}
+
 	for {
 		// If there are no more results, get more now.
 		if len(results) == 0 {
@@ -626,8 +641,10 @@ func (vi *Index) searchForInsertHelper(
 				return nil, err
 			}
 			original = idxCtx.tempSearchSet.PopResults()
-			if len(original) < 1 {
-				return nil, errors.AssertionFailedf("SearchForInsert should return at least one result")
+			if len(original) == 0 {
+				// No results, so try again with more results at each search level.
+				maxResults *= 2
+				continue
 			}
 			results = original
 			allowRetry = false
@@ -669,18 +686,12 @@ func (vi *Index) searchForInsertHelper(
 					return nil, err
 				}
 			} else {
-				// This is a non-root partition, so check the next partition in the
-				// result list. If there are no more partitions to check, then expand
-				// the search to a wider set of partitions by increasing MaxResults.
-				results = results[1:]
-				if len(results) == 0 {
-					if seen == nil {
-						seen = make(map[PartitionKey]struct{})
-					}
-					seen[partitionKey] = struct{}{}
-					maxResults *= 2
-				}
+				// This is a non-root partition, so try next partition.
+				checkNextPartition(partitionKey)
 			}
+		} else if errors.Is(err, ErrPartitionNotFound) {
+			// This partition does not exist, so try next partition.
+			checkNextPartition(partitionKey)
 		} else {
 			return nil, err
 		}
@@ -740,19 +751,19 @@ func (vi *Index) fallbackOnTargets(
 				"fetching centroids for target partitions %d and %d, for splitting partition %d",
 				state.Target1, state.Target2, partitionKey)
 		}
-		if len(tempResults) != 2 {
-			return nil, errors.AssertionFailedf(
-				"expected to get two centroids for state %s", state.String())
-		}
 
-		// Order by the distance of the centroids from the query vector.
-		dist1 := num32.L2SquaredDistance(vec, tempResults[0].Vector)
-		dist2 := num32.L2SquaredDistance(vec, tempResults[1].Vector)
+		// Calculate the distance of the query vector to the centroids.
+		for i := range tempResults {
+			tempResults[i].QuerySquaredDistance = num32.L2SquaredDistance(vec, tempResults[i].Vector)
+		}
 
-		tempResults[0].QuerySquaredDistance = dist1
-		tempResults[1].QuerySquaredDistance = dist2
+		if len(tempResults) < 2 {
+			// One or both of the target partitions have been deleted.
+			return tempResults, nil
+		}
 
-		if dist1 > dist2 {
+		// Order by the distance of the centroids from the query vector.
+		if tempResults[0].QuerySquaredDistance > tempResults[1].QuerySquaredDistance {
 			// Swap results.
 			tempResult := tempResults[0]
 			tempResults[0] = tempResults[1]
@@ -848,12 +859,6 @@ func (vi *Index) searchHelper(ctx context.Context, idxCtx *Context, searchSet *S
 
 	for {
 		results := subSearchSet.PopResults()
-		if len(results) == 0 && searchLevel > LeafLevel {
-			// This should never happen, as it means that interior partition(s)
-			// have no children. The vector deletion logic should prevent that.
-			panic(errors.AssertionFailedf(
-				"interior partition(s) on level %d has no children", searchLevel))
-		}
 
 		var zscore float64
 		if searchLevel > LeafLevel {
@@ -974,13 +979,17 @@ func (vi *Index) searchChildPartitions(
 		}
 	}
 
-	level, err = idxCtx.txn.SearchPartitions(
+	err = idxCtx.txn.SearchPartitions(
 		ctx, idxCtx.treeKey, idxCtx.tempToSearch, idxCtx.randomized, searchSet)
 	if err != nil {
 		return 0, err
 	}
 
 	for i := range parentResults {
+		if level == InvalidLevel {
+			level = idxCtx.tempToSearch[i].Level
+		}
+
 		count := idxCtx.tempToSearch[i].Count
 		searchSet.Stats.SearchedPartition(level, count)
 
@@ -1130,11 +1139,15 @@ func (vi *Index) getFullVectors(
 	for i < len(candidates) {
 		candidates[i].Vector = idxCtx.tempVectorsWithKeys[i].Vector
 
-		// Exclude deleted vectors from results.
+		// Exclude deleted child keys from results.
 		if candidates[i].Vector == nil {
-			// Vector was deleted, so add fixup to delete it.
-			vi.fixups.AddDeleteVector(
-				ctx, candidates[i].ParentPartitionKey, candidates[i].ChildKey.KeyBytes)
+			// TODO(andyk): Need to create an DeletePartitionKey fixup to handle
+			// the case of a dangling partition key.
+			if candidates[i].ChildKey.KeyBytes != nil {
+				// Vector was deleted, so add fixup to delete it.
+				vi.fixups.AddDeleteVector(
+					ctx, candidates[i].ParentPartitionKey, candidates[i].ChildKey.KeyBytes)
+			}
 
 			// Move the last candidate to the current position and reduce size
 			// of slice by one.

pkg/sql/vecindex/cspann/memstore/memstore_test.go

@@ -129,7 +129,7 @@ func TestInMemoryStoreConcurrency(t *testing.T) {
 			commontest.RunTransaction(ctx, t, store, func(txn2 cspann.Txn) {
 				searchSet := cspann.SearchSet{MaxResults: 1}
 				toSearch := []cspann.PartitionToSearch{{Key: cspann.RootKey}}
-				_, err := txn2.SearchPartitions(ctx, treeKey, toSearch, vector.T{0, 0}, &searchSet)
+				err := txn2.SearchPartitions(ctx, treeKey, toSearch, vector.T{0, 0}, &searchSet)
 				require.NoError(t, err)
 				result1 := cspann.SearchResult{
 					QuerySquaredDistance: 25, ErrorBound: 0, CentroidDistance: 5,

pkg/sql/vecindex/cspann/memstore/memstore_txn.go

@@ -231,19 +231,21 @@ func (tx *memTxn) SearchPartitions(
 	toSearch []cspann.PartitionToSearch,
 	queryVector vector.T,
 	searchSet *cspann.SearchSet,
-) (level cspann.Level, err error) {
+) error {
 	for i := range toSearch {
-		var searchLevel cspann.Level
-
 		memPart, ok := tx.store.getPartition(treeKey, toSearch[i].Key)
 		if !ok {
 			if toSearch[i].Key == cspann.RootKey {
 				// Root partition has not yet been created, so it must be empty.
-				searchLevel = cspann.LeafLevel
+				toSearch[i].Level = cspann.LeafLevel
 				toSearch[i].StateDetails = cspann.MakeReadyDetails()
 				toSearch[i].Count = 0
 			} else {
-				return cspann.InvalidLevel, cspann.ErrPartitionNotFound
+				// Partition does not exist, so return InvalidLevel, MissingState
+				// and Count=0.
+				toSearch[i].Level = cspann.InvalidLevel
+				toSearch[i].StateDetails = cspann.PartitionStateDetails{}
+				toSearch[i].Count = 0
 			}
 		} else {
 			// Acquire shared lock on partition and search it. Note that we don't
@@ -255,23 +257,14 @@ func (tx *memTxn) SearchPartitions(
 				defer memPart.lock.ReleaseShared()
 
 				partition := memPart.lock.partition
+				toSearch[i].Level = partition.Level()
 				toSearch[i].StateDetails = partition.Metadata().StateDetails
-				searchLevel, toSearch[i].Count = partition.Search(
-					&tx.workspace, toSearch[i].Key, queryVector, searchSet)
+				toSearch[i].Count = partition.Search(&tx.workspace, toSearch[i].Key, queryVector, searchSet)
 			}()
 		}
-
-		if i == 0 {
-			level = searchLevel
-		} else if level != searchLevel {
-			// Callers should only search for partitions at the same level.
-			panic(errors.AssertionFailedf(
-				"caller already searched a partition at level %d, cannot search at level %d",
-				level, searchLevel))
-		}
 	}
 
-	return level, nil
+	return nil
 }
 
 // GetFullVectors implements the Txn interface.
@@ -286,14 +279,13 @@ func (tx *memTxn) GetFullVectors(
 		if ref.Key.PartitionKey != cspann.InvalidKey {
 			// Get the partition's centroid.
 			memPart, ok := tx.store.getPartitionLocked(treeKey, ref.Key.PartitionKey)
-			if !ok {
-				return errors.Wrapf(cspann.ErrPartitionNotFound,
-					"getting partition %d centroid", ref.Key.PartitionKey)
+			if ok {
+				// Don't need to acquire lock to call the Centroid method, since it
+				// is immutable and thread-safe.
+				ref.Vector = memPart.lock.partition.Centroid()
+			} else {
+				ref.Vector = nil
 			}
-
-			// Don't need to acquire lock to call the Centroid method, since it
-			// is immutable and thread-safe.
-			ref.Vector = memPart.lock.partition.Centroid()
 		} else {
 			vector, ok := tx.store.mu.vectors[string(refs[i].Key.KeyBytes)]
 			if ok {

pkg/sql/vecindex/cspann/partition.go

@@ -130,13 +130,12 @@ func (p *Partition) ValueBytes() []ValueBytes {
 }
 
 // Search estimates the set of data vectors that are nearest to the given query
-// vector and returns them in the given search set. Search also returns this
-// partition's level in the K-means tree and the count of quantized vectors in
-// the partition.
+// vector and returns them in the given search set. Search also returns the
+// count of quantized vectors in the partition.
 func (p *Partition) Search(
 	w *workspace.T, partitionKey PartitionKey, queryVector vector.T, searchSet *SearchSet,
-) (level Level, count int) {
-	count = p.Count()
+) int {
+	count := p.Count()
 	tempFloats := w.AllocFloats(count * 2)
 	defer w.FreeFloats(tempFloats)
 
@@ -161,7 +160,7 @@ func (p *Partition) Search(
 		searchSet.Add(&searchSet.tempResult)
 	}
 
-	return p.Level(), count
+	return count
 }
 
 // Add quantizes the given vector as part of this partition. If a vector with