vecindex: loosen RemoveFromPartition semantics

Previously, vectors could not be removed from partitions that were
in a draining or deleting state. While this restriction is needed
for adding vectors to prevent "lost update" anomalies, it's not
needed for removing vectors. A future PR will be able to avoid extra
work if this restriction is loosened.

Even though it's OK to remove vectors from partitions in any state,
that doesn't mean it's *useful* to do so. For example, when deleting
a vector from the index, there's no point in deleting it from a
partition that is going to be cleared or deleted anyway.

Distinguish between these various cases by splitting AllowAddOrRemove
into AllowAdd, AllowRemove, and CanSkipRemove.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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

@@ -137,10 +137,10 @@ func (suite *StoreTestSuite) TestGetPartitionMetadata() {
 			CheckPartitionMetadata(suite.T(), metadata, cspann.LeafLevel, vector.T{0, 0},
 				cspann.PartitionStateDetails{State: cspann.ReadyState})
 
-			// Non-root partition does not yet exist, expect error.
-			_, err = txn.GetPartitionMetadata(
+			// Non-root partition does not yet exist, expect Missing metadata.
+			metadata, err = txn.GetPartitionMetadata(
 				suite.ctx, treeKey, cspann.PartitionKey(99), false /* forUpdate */)
-			suite.ErrorIs(err, cspann.ErrPartitionNotFound)
+			suite.Equal(cspann.PartitionMetadata{}, metadata)
 		})
 
 		// Create non-root partition with some vectors in it.
@@ -160,21 +160,19 @@ func (suite *StoreTestSuite) TestGetPartitionMetadata() {
 		suite.NoError(store.TryUpdatePartitionMetadata(
 			suite.ctx, treeKey, partitionKey, metadata, expected))
 
+		// Ensure latest metadata gets returned.
 		RunTransaction(suite.ctx, suite.T(), store, func(txn cspann.Txn) {
-			// If forUpdate = false, GetPartitionMetadata should not error.
 			metadata, err := txn.GetPartitionMetadata(
 				suite.ctx, treeKey, partitionKey, false /* forUpdate */)
 			suite.NoError(err)
 			details := cspann.PartitionStateDetails{
 				State: cspann.DrainingForSplitState, Target1: 20, Target2: 30}
 			CheckPartitionMetadata(suite.T(), metadata, cspann.SecondLevel, vector.T{4, 3}, details)
 
-			// If forUpdate = true, GetPartitionMetadata should error.
-			var errConditionFailed *cspann.ConditionFailedError
-			_, err = txn.GetPartitionMetadata(suite.ctx, treeKey, partitionKey, true /* forUpdate */)
-			suite.ErrorAs(err, &errConditionFailed)
-			CheckPartitionMetadata(suite.T(), errConditionFailed.Actual, cspann.SecondLevel,
-				vector.T{4, 3}, details)
+			metadata, err = txn.GetPartitionMetadata(
+				suite.ctx, treeKey, partitionKey, true /* forUpdate */)
+			suite.NoError(err)
+			CheckPartitionMetadata(suite.T(), metadata, cspann.SecondLevel, vector.T{4, 3}, details)
 		})
 		suite.NoError(err)
 	}
@@ -329,16 +327,12 @@ func (suite *StoreTestSuite) TestRemoveFromPartition() {
 			suite.ctx, treeKey, partitionKey, metadata, expected))
 
 		RunTransaction(suite.ctx, suite.T(), store, func(txn cspann.Txn) {
-			// Try to remove from partition, expect error due to its state.
-			var errConditionFailed *cspann.ConditionFailedError
+			// Try to remove from draining partition, expect success.
 			err := txn.RemoveFromPartition(suite.ctx, treeKey, partitionKey, cspann.SecondLevel,
 				partitionKey3)
-			suite.ErrorAs(err, &errConditionFailed)
-			details := cspann.PartitionStateDetails{
-				State: cspann.DrainingForSplitState, Target1: 20, Target2: 30}
-			CheckPartitionMetadata(suite.T(), errConditionFailed.Actual, cspann.SecondLevel,
-				vector.T{4, 3}, details)
+			suite.NoError(err)
 		})
+		CheckPartitionCount(suite.ctx, suite.T(), store, treeKey, partitionKey, 1)
 	}
 
 	suite.Run("default tree", func() {
@@ -1004,9 +998,27 @@ func (suite *StoreTestSuite) TestTryClearPartition() {
 		// Create partition with some vectors.
 		partitionKey, partition := suite.createTestPartition(store, treeKey)
 
-		// Now clear should work.
+		// Clear should fail in Ready state.
 		expected := *partition.Metadata()
 		count, err := store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
+		suite.Error(err)
+		suite.Equal(0, count)
+
+		// Move to draining state.
+		metadata := *partition.Metadata()
+		metadata.StateDetails.MakeDrainingForSplit(20, 30)
+		suite.NoError(store.TryUpdatePartitionMetadata(
+			suite.ctx, treeKey, partitionKey, metadata, expected))
+
+		// Try to clear with mismatched expected metadata.
+		var errConditionFailed *cspann.ConditionFailedError
+		_, err = store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
+		suite.ErrorAs(err, &errConditionFailed)
+		suite.True(errConditionFailed.Actual.Equal(&metadata))
+
+		// Try again, this time with correct expected metadata.
+		expected = metadata
+		count, err = store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
 		suite.NoError(err)
 		suite.Equal(3, count)
 
@@ -1019,15 +1031,7 @@ func (suite *StoreTestSuite) TestTryClearPartition() {
 		suite.Len(partition.ChildKeys(), 0)
 		suite.Len(partition.ValueBytes(), 0)
 
-		// 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.
+		// Clear partition with zero vectors.
 		count, err = store.TryClearPartition(suite.ctx, treeKey, partitionKey, expected)
 		suite.NoError(err)
 		suite.Equal(0, count)

pkg/sql/vecindex/cspann/fixup_split.go

@@ -299,7 +299,7 @@ func (fw *fixupWorker) splitPartition(
 			}
 		}
 
-		// Remove the splitting partition from the its parent. Note that we don't
+		// Remove the splitting partition from its parent. Note that we don't
 		// delete the partition's metadata record, instead leaving it behind as a
 		// "tombstone". This prevents other racing workers from resurrecting the
 		// partition as a zombie, which could otherwise happen like this:
@@ -457,7 +457,7 @@ func (fw *fixupWorker) reassignToSiblings(
 		for offset, distance := range tempSiblingDistances {
 			if distance >= minDistance {
 				continue
-			} else if !fw.tempMetadataToGet[offset].Metadata.StateDetails.State.AllowAddOrRemove() {
+			} else if !fw.tempMetadataToGet[offset].Metadata.StateDetails.State.AllowAdd() {
 				continue
 			}
 			siblingOffset = offset
@@ -593,9 +593,9 @@ func (fw *fixupWorker) addToPartition(
 	valueBytes []ValueBytes,
 	expected PartitionMetadata,
 ) (added bool, err error) {
-	if !expected.StateDetails.State.AllowAddOrRemove() {
+	if !expected.StateDetails.State.AllowAdd() {
 		return false, errors.AssertionFailedf(
-			"cannot add to partition in state that disallows adds/removes")
+			"cannot add to partition in state %s that disallows adds", expected.StateDetails.State)
 	}
 	fw.index.validateVectorsToAdd(expected.Level, vectors)
 
@@ -623,10 +623,11 @@ func (fw *fixupWorker) addToPartition(
 func (fw *fixupWorker) clearPartition(
 	ctx context.Context, partitionKey PartitionKey, metadata PartitionMetadata,
 ) (err error) {
-	if metadata.StateDetails.State.AllowAddOrRemove() {
+	if metadata.StateDetails.State.AllowAdd() {
+		// Something's wrong if partition is being cleared in a state that allows
+		// new vectors to be added.
 		return errors.AssertionFailedf(
-			"cannot clear partition %d in state %s that allows adds/removes",
-			partitionKey, metadata.StateDetails.String())
+			"cannot add to partition in state %s that allows adds", metadata.StateDetails.State)
 	}
 
 	// Remove all children in the partition.
@@ -823,10 +824,9 @@ func (fw *fixupWorker) removeFromParentPartition(
 		return errors.Wrapf(err, "getting parent partition %d metadata", parentPartitionKey)
 	}
 
-	if !parentMetadata.StateDetails.State.AllowAddOrRemove() || parentMetadata.Level != parentLevel {
-		// Child could not be removed from the parent because it doesn't exist or
-		// it no longer allows deletes, or its level has changed (i.e. in root
-		// partition case).
+	if parentMetadata.StateDetails.State.CanSkipRemove() || parentMetadata.Level != parentLevel {
+		// Child should not be removed from the parent that's draining or will be
+		// deleted, or if its level has changed (i.e. in root partition case).
 		return errFixupAborted
 	}
 

pkg/sql/vecindex/cspann/index.go

@@ -504,6 +504,10 @@ func (vi *Index) SearchForInsert(
 		if err != nil {
 			return errors.Wrapf(err, "locking metadata for insert into partition %d", partitionKey)
 		}
+		if !metadata.StateDetails.State.AllowAdd() {
+			// The partition does not allow adds, so go to the next candidate.
+			return NewConditionFailedError(metadata)
+		}
 		result.Vector = metadata.Centroid
 		return nil
 	}
@@ -533,10 +537,16 @@ func (vi *Index) SearchForDelete(
 	// When a candidate delete partition is found, lock its metadata for update.
 	removeFunc := func(ctx context.Context, idxCtx *Context, result *SearchResult) error {
 		partitionKey := result.ParentPartitionKey
-		_, err := idxCtx.txn.GetPartitionMetadata(ctx, treeKey, partitionKey, true /* forUpdate */)
+		metadata, err := idxCtx.txn.GetPartitionMetadata(
+			ctx, treeKey, partitionKey, true /* forUpdate */)
 		if err != nil {
 			return errors.Wrapf(err, "locking metadata for delete from partition %d", partitionKey)
 		}
+		if metadata.StateDetails.State.CanSkipRemove() {
+			// The partition will be cleared or deleted anyway, so no need to
+			// remove from it. Go to the next candidate.
+			return NewConditionFailedError(metadata)
+		}
 		return nil
 	}
 
@@ -685,6 +695,10 @@ func (vi *Index) searchForUpdateHelper(
 			}
 			if !ok {
 				if idxCtx.forInsert {
+					// Keep aggressively searching for valid insert partition, since
+					// the only alternative is to fail the operation. This is not
+					// necessary in the delete case, since it's OK if there are
+					// dangling vectors in rare cases.
 					return vi.searchForUpdateHelper(ctx, idxCtx, fn, deleteKey, remainingAttempts)
 				}
 				break

pkg/sql/vecindex/cspann/index_test.go

@@ -818,11 +818,12 @@ func (s *testState) loadIndexFromFormat(
 		}
 	}
 
+	// Always create partition in Ready state so that adds are allowed.
 	metadata := cspann.PartitionMetadata{
-		Level:        childLevel,
-		Centroid:     centroid,
-		StateDetails: details,
+		Level:    childLevel,
+		Centroid: centroid,
 	}
+	metadata.StateDetails.MakeReady()
 	err = s.MemStore.TryCreateEmptyPartition(s.Ctx, treeKey, partitionKey, metadata)
 	require.NoError(s.T, err)
 
@@ -834,6 +835,14 @@ func (s *testState) loadIndexFromFormat(
 		require.True(s.T, added)
 	}
 
+	if details.State != cspann.ReadyState {
+		// Update the partition's state
+		expected := metadata
+		metadata.StateDetails = details
+		err = s.MemStore.TryUpdatePartitionMetadata(s.Ctx, treeKey, partitionKey, metadata, expected)
+		require.NoError(s.T, err)
+	}
+
 	return lines, childLevel + 1, centroid, cspann.ChildKey{PartitionKey: partitionKey}
 }
 

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

@@ -467,6 +467,10 @@ func (s *Store) TryAddToPartition(
 	if !existing.Equal(&expected) {
 		return false, cspann.NewConditionFailedError(*existing)
 	}
+	if !existing.StateDetails.State.AllowAdd() {
+		return false, errors.AssertionFailedf(
+			"cannot add to partition in state %s that disallows adds", existing.StateDetails.State)
+	}
 
 	// Add the vectors to the partition. Ignore any duplicate vectors.
 	// TODO(andyk): Figure out how to give Store flexible scratch space.
@@ -520,15 +524,19 @@ func (s *Store) TryClearPartition(
 	memPart := s.lockPartition(treeKey, partitionKey, uniqueOwner, true /* isExclusive */)
 	if memPart == nil {
 		// Partition does not exist.
-		return -1, cspann.ErrPartitionNotFound
+		return 0, cspann.ErrPartitionNotFound
 	}
 	defer memPart.lock.Release()
 
 	// Check precondition.
 	partition := memPart.lock.partition
 	existing := partition.Metadata()
 	if !existing.Equal(&expected) {
-		return -1, cspann.NewConditionFailedError(*existing)
+		return 0, cspann.NewConditionFailedError(*existing)
+	}
+	if existing.StateDetails.State.AllowAdd() {
+		return 0, errors.AssertionFailedf(
+			"cannot clear partition in state %s that allows adds", existing.StateDetails.State)
 	}
 
 	// Remove vectors from the partition and update partition count.

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

@@ -79,16 +79,7 @@ func (tx *memTxn) GetPartitionMetadata(
 		defer memPart.lock.ReleaseShared()
 	}
 
-	// Do not allow updates to the partition if the state doesn't allow it.
-	metadata := memPart.lock.partition.Metadata()
-	if forUpdate && !metadata.StateDetails.State.AllowAddOrRemove() {
-		err = cspann.NewConditionFailedError(*metadata)
-		return cspann.PartitionMetadata{}, errors.Wrapf(err,
-			"getting metadata for partition %d (state=%s)",
-			partitionKey, metadata.StateDetails.State.String())
-	}
-
-	return *metadata, nil
+	return *memPart.lock.partition.Metadata(), nil
 }
 
 // AddToPartition implements the Txn interface.
@@ -117,7 +108,7 @@ func (tx *memTxn) AddToPartition(
 	// allow it.
 	partition := memPart.lock.partition
 	state := partition.Metadata().StateDetails.State
-	if !state.AllowAddOrRemove() {
+	if !state.AllowAdd() {
 		return errors.Wrapf(cspann.NewConditionFailedError(*partition.Metadata()),
 			"adding to partition %d (state=%s)", partitionKey, state.String())
 	}
@@ -158,16 +149,8 @@ func (tx *memTxn) RemoveFromPartition(
 	}
 	defer memPart.lock.Release()
 
-	// Do not allow vectors to be removed from the partition if the state doesn't
-	// allow it.
-	partition := memPart.lock.partition
-	state := partition.Metadata().StateDetails.State
-	if !state.AllowAddOrRemove() {
-		return errors.Wrapf(cspann.NewConditionFailedError(*partition.Metadata()),
-			"removing from partition %d (state=%s)", partitionKey, state.String())
-	}
-
 	// Remove vector from the partition.
+	partition := memPart.lock.partition
 	if level != partition.Level() {
 		return errors.Wrapf(cspann.ErrRestartOperation,
 			"removing from partition %d (expected: %d, actual: %d)",

pkg/sql/vecindex/cspann/partition_metadata.go

@@ -45,9 +45,8 @@ func ParsePartitionState(s string) PartitionState {
 	return MissingState
 }
 
-// AllowAddOrRemove returns true if vectors can be added or removed to/from the
-// partition in this state.
-func (s PartitionState) AllowAddOrRemove() bool {
+// AllowAdd returns true if vectors can be added to the partition in this state.
+func (s PartitionState) AllowAdd() bool {
 	switch s {
 	case ReadyState, SplittingState, MergingState, UpdatingState,
 		AddingLevelState, RemovingLevelState:
@@ -56,6 +55,18 @@ func (s PartitionState) AllowAddOrRemove() bool {
 	return false
 }
 
+// CanSkipRemove returns true if there's no need to remove vectors from the
+// partition because the partition is being drained or is about to be deleted.
+func (s PartitionState) CanSkipRemove() bool {
+	return !s.AllowAdd()
+}
+
+// AllowRemove returns true if vectors can be removed from the partition in this
+// state.
+func (s PartitionState) AllowRemove() bool {
+	return s != MissingState
+}
+
 // String returns the state formatted as a string.
 func (s PartitionState) String() string {
 	switch s {
@@ -104,12 +115,12 @@ const (
 	UpdatingState
 	// DrainingForSplitState indicates that the partition is actively moving
 	// vectors to target split sub-partitions. Searches are allowed, but not
-	// inserts, deletes, splits, or merges.
+	// inserts, splits, or merges.
 	DrainingForSplitState
 	// DrainingForMergeState indicates that the partition is actively moving
 	// vectors into other partitions at the same level (or deleting vectors if
-	// this is the root partition). Searches are allowed, but not inserts,
-	// deletes, splits, or merges.
+	// this is the root partition). Searches are allowed, but not inserts, splits,
+	// or merges.
 	DrainingForMergeState
 	// AddingLevelState indicates that a root partition has been drained after a
 	// split and has had its level increased by one. What remains is to add the
@@ -123,7 +134,7 @@ const (
 	RemovingLevelState
 	// DeletingForSplitState indicates that a non-root splitting partition is
 	// about to be removed from its tree and deleted. Searches are allowed, but
-	// not inserts, deletes, splits, or merges.
+	// not inserts, splits, or merges.
 	DeletingForSplitState
 )
 

pkg/sql/vecindex/cspann/searcher.go

@@ -422,16 +422,16 @@ func (s *levelSearcher) searchChildPartitions(
 		}
 		s.stats.SearchedPartition(level, count)
 
-		// If searching for vector to delete, skip partitions that are in a state
-		// that does not allow add and remove operations. This is not possible to
-		// do here for the insert case, because we do not actually search the
-		// partition in which to insert; we only search its parent and never get
-		// the metadata for the insert partition itself.
+		// If searching for vector to delete, skip partitions that don't need
+		// vectors deleted from them (because they are draining or deleting). This
+		// is not possible to do here for the insert case, because we do not
+		// actually search the partition in which to insert; we only search its
+		// parent and never get the metadata for the insert partition itself.
 		// TODO(andyk): This should probably be checked in the Store, perhaps by
 		// passing a "forUpdate" parameter to SearchPartitions, so that the Store
 		// doesn't even add vectors from partitions that do not allow updates.
 		if s.idxCtx.forDelete && s.idxCtx.level == level {
-			if !s.idxCtx.tempToSearch[i].StateDetails.State.AllowAddOrRemove() {
+			if s.idxCtx.tempToSearch[i].StateDetails.State.CanSkipRemove() {
 				s.searchSet.RemoveByParent(partitionKey)
 			}
 		}