cspann: make vector delete work during incremental splits

The new split code performs splits of vector index K-means trees using an
incremental series of non-transactional steps. This commit ensures that
any deletes from the vector index work as expected during all of these steps.
In particular, in the DrainingForSplit phase, a partition does not allow
vectors to be deleted. Instead, the vector needs to be deleted from one of
the target partitions of the split (i.e. the two partitions between which
the vectors are divided).

Epic: CRDB-42943

Release note: None

Author: andy-kimball

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

@@ -1094,6 +1094,64 @@ func (suite *StoreTestSuite) TestTryRemoveFromPartition() {
 	}
 }
 
+func (suite *StoreTestSuite) TestTryClearPartition() {
+	store := suite.makeStore(suite.quantizer)
+	if !store.SupportsTry() {
+		return
+	}
+
+	doTest := func(treeID int) {
+		treeKey := store.MakeTreeKey(suite.T(), treeID)
+
+		// Partition does not yet exist.
+		_, err := store.TryClearPartition(suite.ctx, treeKey, cspann.PartitionKey(99),
+			cspann.PartitionMetadata{})
+		suite.ErrorIs(err, cspann.ErrPartitionNotFound)
+
+		// Create partition with some vectors.
+		partitionKey, partition := suite.createTestPartition(store, treeKey)
+
+		// Now clear should work.
+		expected := *partition.Metadata()
+		count, err := store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
+		suite.NoError(err)
+		suite.Equal(3, count)
+
+		// Fetch back the partition and validate it.
+		partition, err = store.TryGetPartition(suite.ctx, treeKey, partitionKey)
+		suite.NoError(err)
+		suite.True(partition.Metadata().Equal(&expected))
+		suite.Equal(cspann.SecondLevel, partition.Level())
+		suite.Equal(vector.T{4, 3}, partition.Centroid())
+		suite.Equal([]cspann.ChildKey{}, partition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{}, partition.ValueBytes())
+
+		// Try to clear with mismatched expected metadata.
+		var errConditionFailed *cspann.ConditionFailedError
+		metadata := expected
+		metadata.StateDetails.State = cspann.DrainingForMergeState
+		_, err = store.TryClearPartition(suite.ctx, treeKey, partitionKey, metadata)
+		suite.ErrorAs(err, &errConditionFailed)
+		suite.True(errConditionFailed.Actual.Equal(&expected))
+
+		// Try again, this time with correct expected metadata.
+		count, err = store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
+		suite.NoError(err)
+		suite.Equal(0, count)
+	}
+
+	suite.Run("default tree", func() {
+		doTest(0)
+	})
+
+	if store.AllowMultipleTrees() {
+		// Ensure that vectors are independent across trees.
+		suite.Run("different tree", func() {
+			doTest(1)
+		})
+	}
+}
+
 func (suite *StoreTestSuite) runInTransaction(store TestStore, fn func(tx cspann.Txn)) {
 	suite.NoError(store.RunTransaction(suite.ctx, func(tx cspann.Txn) error {
 		fn(tx)

pkg/sql/vecindex/cspann/fixup_split.go

@@ -245,7 +245,7 @@ func (fw *fixupWorker) splitPartition(
 			// This is the root partition, so remove all of its vectors rather than
 			// delete the root partition itself. Note that the vectors have already
 			// been copied to the two target partitions.
-			err = fw.clearPartition(ctx, partitionKey, partition)
+			err = fw.clearPartition(ctx, partitionKey, *partition.Metadata())
 			if err != nil {
 				return err
 			}
@@ -374,29 +374,28 @@ func (fw *fixupWorker) addToPartition(
 	return nil
 }
 
-// clearPartition removes all vectors and associated data from the given
-// partition, leaving it empty, on the condition that the partition's state has
-// not changed unexpectedly. If that's the case, it returns errFixupAborted.
+// clearPartition removes all vectors from the given partition. This only
+// happens if the partition's metadata has not changed. If it has changed,
+// clearPartition returns errFixupAborted.
 func (fw *fixupWorker) clearPartition(
-	ctx context.Context, partitionKey PartitionKey, partition *Partition,
+	ctx context.Context, partitionKey PartitionKey, metadata PartitionMetadata,
 ) (err error) {
-	if partition.Metadata().StateDetails.State.AllowAddOrRemove() {
+	if metadata.StateDetails.State.AllowAddOrRemove() {
 		return errors.AssertionFailedf("cannot clear partition in state that allows adds/removes")
 	}
 
 	// Remove all children in the partition.
-	removed, err := fw.index.store.TryRemoveFromPartition(ctx, fw.treeKey,
-		partitionKey, partition.ChildKeys(), *partition.Metadata())
+	count, err := fw.index.store.TryClearPartition(ctx, fw.treeKey, partitionKey, metadata)
 	if err != nil {
 		metadata, err := suppressRaceErrors(err)
 		if err == nil {
 			// Another worker raced to update the metadata, so abort.
 			return errors.Wrapf(errFixupAborted,
-				"clearing % vectors from partition, %d expected %s, found %s", partition.Count(),
+				"clearing vectors from partition %d, expected %s, found %s",
 				partitionKey, metadata.StateDetails.String(), metadata.StateDetails.String())
 		}
 		return errors.Wrap(err, "clearing vectors")
-	} else if fw.singleStep && removed {
+	} else if fw.singleStep && count > 0 {
 		return errFixupAborted
 	}
 
@@ -549,7 +548,7 @@ func (fw *fixupWorker) addToParentPartition(
 func (fw *fixupWorker) deletePartition(
 	ctx context.Context, parentPartitionKey, partitionKey PartitionKey,
 ) (err error) {
-	const format = "deleting partition %d, with parent partition %d (state=%d)"
+	const format = "deleting partition %d (parent=%d, state=%d)"
 	var parentMetadata PartitionMetadata
 
 	defer func() {

pkg/sql/vecindex/cspann/index.go

@@ -131,6 +131,9 @@ type Context struct {
 	// forInsert indicates that this is an insert operation (or a search for
 	// insert operation).
 	forInsert bool
+	// forDelete indicates that this is a delete operation (or a search for
+	// delete operation).
+	forDelete bool
 
 	tempSearchSet       SearchSet
 	tempSubSearchSet    SearchSet
@@ -449,14 +452,15 @@ func (vi *Index) SearchForDelete(
 		SkipRerank:  true,
 		UpdateStats: true,
 	}, LeafLevel)
+	idxCtx.forDelete = true
 
 	idxCtx.tempSearchSet = SearchSet{MaxResults: 1, MatchKey: key}
 
 	// Search with the base beam size. If that fails to find the vector, try again
 	// with a larger beam size, in order to minimize the chance of dangling
 	// vector references in the index.
 	baseBeamSize := max(vi.options.BaseBeamSize, 1)
-	for i := 0; i < 2; i++ {
+	for range 2 {
 		idxCtx.options.BaseBeamSize = baseBeamSize
 
 		err := vi.searchHelper(ctx, idxCtx, &idxCtx.tempSearchSet)
@@ -792,7 +796,7 @@ func (vi *Index) searchHelper(ctx context.Context, idxCtx *Context, searchSet *S
 			// Compute the Z-score of the candidate list if there are enough
 			// samples. Otherwise, use the default Z-score of 0.
 			if len(results) >= vi.options.QualitySamples {
-				for i := 0; i < vi.options.QualitySamples; i++ {
+				for i := range vi.options.QualitySamples {
 					idxCtx.tempQualitySamples[i] = float64(results[i].QuerySquaredDistance)
 				}
 				samples := idxCtx.tempQualitySamples[:vi.options.QualitySamples]
@@ -923,6 +927,14 @@ func (vi *Index) searchChildPartitions(
 			searchSet.RemoveByParent(partitionKey)
 		}
 
+		// If searching for vector to delete, skip partitions that are in a state
+		// that does not allow add and remove operations.
+		if idxCtx.forDelete && idxCtx.level == level {
+			if !idxCtx.tempToSearch[i].StateDetails.State.AllowAddOrRemove() {
+				searchSet.RemoveByParent(partitionKey)
+			}
+		}
+
 		// Enqueue background fixup if a split or merge operation needs to be
 		// started or continued after stalling.
 		state := idxCtx.tempToSearch[i].StateDetails
@@ -1050,7 +1062,8 @@ func (vi *Index) getFullVectors(
 		return nil, err
 	}
 
-	for i := 0; i < len(candidates); i++ {
+	i := 0
+	for i < len(candidates) {
 		candidates[i].Vector = idxCtx.tempVectorsWithKeys[i].Vector
 
 		// Exclude deleted vectors from results.
@@ -1065,7 +1078,8 @@ func (vi *Index) getFullVectors(
 			candidates[i] = candidates[len(candidates)-1]
 			candidates[len(candidates)-1] = SearchResult{} // for GC
 			candidates = candidates[:len(candidates)-1]
-			i--
+		} else {
+			i++
 		}
 	}
 

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

@@ -544,6 +544,34 @@ func (s *Store) TryRemoveFromPartition(
 	return removed, err
 }
 
+// TryClearPartition implements the Store interface.
+func (s *Store) TryClearPartition(
+	ctx context.Context,
+	treeKey cspann.TreeKey,
+	partitionKey cspann.PartitionKey,
+	expected cspann.PartitionMetadata,
+) (count int, err error) {
+	memPart := s.lockPartition(treeKey, partitionKey, uniqueOwner, true /* isExclusive */)
+	if memPart == nil {
+		// Partition does not exist.
+		return -1, cspann.ErrPartitionNotFound
+	}
+	defer memPart.lock.Release()
+
+	// Check precondition.
+	partition := memPart.lock.partition
+	existing := partition.Metadata()
+	if !existing.Equal(&expected) {
+		return -1, cspann.NewConditionFailedError(*existing)
+	}
+
+	// Remove vectors from the partition and update partition count.
+	count = partition.Clear()
+	memPart.count.Store(0)
+
+	return count, nil
+}
+
 // EnsureUniquePartitionKey checks that the given partition key is not being
 // used yet and also ensures it won't be given out in the future. This is used
 // for testing.

pkg/sql/vecindex/cspann/partition.go

@@ -253,6 +253,18 @@ func (p *Partition) Find(childKey ChildKey) int {
 	return -1
 }
 
+// Clear removes all vectors from the partition and returns the number of
+// vectors that were cleared. The centroid stays the same.
+func (p *Partition) Clear() int {
+	count := len(p.childKeys)
+	p.quantizedSet.Clear(p.quantizedSet.GetCentroid())
+	clear(p.childKeys)
+	p.childKeys = p.childKeys[:0]
+	clear(p.valueBytes)
+	p.valueBytes = p.valueBytes[:0]
+	return count
+}
+
 // CreateEmptyPartition returns an empty partition for the given quantizer and
 // level.
 func CreateEmptyPartition(quantizer quantize.Quantizer, metadata PartitionMetadata) *Partition {

pkg/sql/vecindex/cspann/partition_test.go

@@ -237,6 +237,22 @@ func TestPartition(t *testing.T) {
 		require.Equal(t, -1, partition.Find(childKey40))
 		require.False(t, partition.ReplaceWithLastByKey(childKey40))
 	})
+
+	t.Run("test Clear", func(t *testing.T) {
+		partition := newTestPartition()
+		require.Equal(t, 3, partition.Clear())
+		require.Equal(t, 0, partition.Count())
+		require.Equal(t, []ChildKey{}, partition.ChildKeys())
+		require.Equal(t, []ValueBytes{}, partition.ValueBytes())
+		require.Equal(t, []float32{4, 3.33}, testutils.RoundFloats(partition.Centroid(), 2))
+
+		// Clear empty partition.
+		require.Equal(t, 0, partition.Clear())
+		require.Equal(t, 0, partition.Count())
+		require.Equal(t, []ChildKey{}, partition.ChildKeys())
+		require.Equal(t, []ValueBytes{}, partition.ValueBytes())
+		require.Equal(t, []float32{4, 3.33}, testutils.RoundFloats(partition.Centroid(), 2))
+	})
 }
 
 func roundResults(results SearchResults, prec int) SearchResults {

pkg/sql/vecindex/cspann/store.go

@@ -163,6 +163,18 @@ type Store interface {
 		childKeys []ChildKey,
 		expected PartitionMetadata,
 	) (removed bool, err error)
+
+	// TryClearPartition removes all vectors in the specified partition and
+	// returns the number of vectors that were cleared. It returns
+	// ErrPartitionNotFound if the partition does not exist.
+	//
+	// Before performing any action, TryClearPartition checks the partition's
+	// metadata and returns a ConditionFailedError if it is not the same as the
+	// expected metadata. If the partition does not exist, it returns
+	// ErrPartitionNotFound.
+	TryClearPartition(
+		ctx context.Context, treeKey TreeKey, partitionKey PartitionKey, expected PartitionMetadata,
+	) (count int, err error)
 }
 
 // Txn enables callers to make changes to the stored index in a transactional