vecindex: add TryMoveVector method to Store interface

Add a new TryMoveVector to the Store interface. This method moves a
single vector from one partition to another as an atomic operation,
provided the destination metadata is as expected.

Epic: CRDB-42943
Release note: None

Author: andy-kimball

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

@@ -6,7 +6,9 @@
 package commontest
 
 import (
+	"cmp"
 	"context"
+	"slices"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -987,6 +989,135 @@ func (suite *StoreTestSuite) TestTryRemoveFromPartition() {
 	}
 }
 
+func (suite *StoreTestSuite) TestTryMoveVector() {
+	store := suite.makeStore(suite.quantizer)
+	defer store.Close(suite.T())
+
+	doTest := func(treeID int) {
+		treeKey := store.MakeTreeKey(suite.T(), treeID)
+
+		// Create source partition with some vectors.
+		sourcePartitionKey, _ := suite.createTestPartition(store, treeKey)
+
+		// Create empty target partition.
+		targetPartitionKey := cspann.PartitionKey(20)
+		metadata := cspann.PartitionMetadata{
+			Level:    cspann.SecondLevel,
+			Centroid: vector.T{2, 4},
+		}
+		metadata.StateDetails.MakeReady()
+		suite.NoError(store.TryCreateEmptyPartition(suite.ctx, treeKey, targetPartitionKey, metadata))
+		targetPartition, err := store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+
+		// Source partition does not yet exist.
+		expected := *targetPartition.Metadata()
+		moved, err := store.TryMoveVector(
+			suite.ctx, treeKey, cspann.PartitionKey(99), targetPartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Destination partition does not yet exist.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, cspann.PartitionKey(99),
+			vec1, partitionKey1, valueBytes1, cspann.PartitionMetadata{})
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Source partition is the same as destination partition.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, sourcePartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Now move should work.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey1, valueBytes1, expected)
+		suite.NoError(err)
+		suite.True(moved)
+
+		// Fetch back the target partition and validate it.
+		targetPartition, err = store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey1}, targetPartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes1}, targetPartition.ValueBytes())
+
+		// Try to move again, but with mismatched expected metadata.
+		var errConditionFailed *cspann.ConditionFailedError
+		metadata = expected
+		metadata.StateDetails.State = cspann.DrainingForMergeState
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, metadata)
+		suite.ErrorAs(err, &errConditionFailed)
+		suite.False(moved)
+		suite.True(errConditionFailed.Actual.Equal(&expected))
+
+		// Try again, this time with correct expected metadata.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, expected)
+		suite.NoError(err)
+		suite.True(moved)
+
+		// Fetch back the source partition and validate it.
+		sourcePartition, err := store.TryGetPartition(suite.ctx, treeKey, sourcePartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey2}, sourcePartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes2}, sourcePartition.ValueBytes())
+
+		// Try to move a vector that no longer exists in the source partition.
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey3, valueBytes3, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Try to move a vector that already exists in the target partition.
+		added, err := store.TryAddToPartition(
+			suite.ctx, treeKey, targetPartitionKey, vec2.AsSet(),
+			[]cspann.ChildKey{partitionKey2}, []cspann.ValueBytes{valueBytes2}, expected)
+		suite.NoError(err)
+		suite.True(added)
+
+		moved, err = store.TryMoveVector(
+			suite.ctx, treeKey, sourcePartitionKey, targetPartitionKey,
+			vec1, partitionKey2, valueBytes2, expected)
+		suite.NoError(err)
+		suite.False(moved)
+
+		// Ensure that the vector was not removed from the source partition.
+		sourcePartition, err = store.TryGetPartition(suite.ctx, treeKey, sourcePartitionKey)
+		suite.NoError(err)
+		suite.Equal([]cspann.ChildKey{partitionKey2}, sourcePartition.ChildKeys())
+		suite.Equal([]cspann.ValueBytes{valueBytes2}, sourcePartition.ValueBytes())
+
+		// Ensure that the target partition now has all three vectors.
+		targetPartition, err = store.TryGetPartition(suite.ctx, treeKey, targetPartitionKey)
+		suite.NoError(err)
+		suite.Equal(3, targetPartition.Count())
+		childKeys := slices.Clone(targetPartition.ChildKeys())
+		slices.SortFunc(childKeys, func(a, b cspann.ChildKey) int {
+			return cmp.Compare(a.PartitionKey, b.PartitionKey)
+		})
+		suite.Equal([]cspann.ChildKey{partitionKey1, partitionKey2, partitionKey3}, childKeys)
+	}
+
+	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) TestTryClearPartition() {
 	store := suite.makeStore(suite.quantizer)
 	defer store.Close(suite.T())

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

@@ -514,6 +514,76 @@ func (s *Store) TryRemoveFromPartition(
 	return removed, err
 }
 
+// TryMoveVector implements the Store interface.
+func (s *Store) TryMoveVector(
+	ctx context.Context,
+	treeKey cspann.TreeKey,
+	sourcePartitionKey, targetPartitionKey cspann.PartitionKey,
+	vec vector.T,
+	childKey cspann.ChildKey,
+	valueBytes cspann.ValueBytes,
+	expected cspann.PartitionMetadata,
+) (moved bool, err error) {
+	if sourcePartitionKey == targetPartitionKey {
+		// No-op move.
+		return false, nil
+	}
+
+	// Always lock the partition with the lower key first, in order to avoid
+	// deadlocks.
+	partitionKey1 := sourcePartitionKey
+	partitionKey2 := targetPartitionKey
+	if partitionKey2 < partitionKey1 {
+		partitionKey1, partitionKey2 = partitionKey2, partitionKey1
+	}
+	sourceMemPart := s.lockPartition(treeKey, partitionKey1, uniqueOwner, true /* isExclusive */)
+	if sourceMemPart == nil {
+		// Partition does not exist.
+		return false, nil
+	}
+	defer sourceMemPart.lock.Release()
+
+	targetMemPart := s.lockPartition(treeKey, partitionKey2, uniqueOwner, true /* isExclusive */)
+	if targetMemPart == nil {
+		// Partition does not exist.
+		return false, nil
+	}
+	defer targetMemPart.lock.Release()
+
+	if targetPartitionKey < sourcePartitionKey {
+		sourceMemPart, targetMemPart = targetMemPart, sourceMemPart
+	}
+
+	// Check precondition against target partition.
+	targetPartition := targetMemPart.lock.partition
+	existing := targetPartition.Metadata()
+	if !existing.Equal(&expected) {
+		return false, cspann.NewConditionFailedError(*existing)
+	}
+
+	// Ensure vector is in the source partition, but don't remove it yet.
+	sourcePartition := sourceMemPart.lock.partition
+	if sourcePartition.Find(childKey) == -1 {
+		return false, nil
+	}
+
+	// Add vector to the target partition.
+	// TODO(andyk): Figure out how to give Store flexible scratch space.
+	var w workspace.T
+	added := targetPartition.Add(&w, vec, childKey, valueBytes, false /* overwrite */)
+	if !added {
+		// Vector already exists in the target partition.
+		return false, nil
+	}
+	targetMemPart.count.Store(int64(targetPartition.Count()))
+
+	// Remove vector from the source partition.
+	_ = sourcePartition.ReplaceWithLastByKey(childKey)
+	sourceMemPart.count.Store(int64(sourcePartition.Count()))
+
+	return true, nil
+}
+
 // TryClearPartition implements the Store interface.
 func (s *Store) TryClearPartition(
 	ctx context.Context,

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

@@ -40,63 +40,63 @@ type memLock struct {
 
 // IsAcquiredBy returns true if the lock is exclusively owned by the given
 // owner, or false if not.
-func (pl *memLock) IsAcquiredBy(owner uint64) bool {
-	return owner != uniqueOwner && pl.exclusiveOwner.Load() == owner
+func (l *memLock) IsAcquiredBy(owner uint64) bool {
+	return owner != uniqueOwner && l.exclusiveOwner.Load() == owner
 }
 
 // Acquire obtains exclusive write access to the resource protected by this
 // lock. The same owner can obtain the lock multiple times. The caller must
 // ensure that Release is called for each call to Acquire.
-func (pl *memLock) Acquire(owner uint64) {
-	if owner != uniqueOwner && pl.exclusiveOwner.Load() == owner {
+func (l *memLock) Acquire(owner uint64) {
+	if owner != uniqueOwner && l.exclusiveOwner.Load() == owner {
 		// Exclusive lock has already been acquired by this owner.
-		pl.mu.reentrancy++
+		l.mu.reentrancy++
 		return
 	}
 
 	// Block until exclusive lock is acquired.
-	pl.mu.Lock()
-	pl.exclusiveOwner.Store(owner) //nolint:deferunlockcheck
+	l.mu.Lock()
+	l.exclusiveOwner.Store(owner) //nolint:deferunlockcheck
 }
 
 // AcquireShared obtains shared read access to the resource protected by this
 // lock. The same owner can obtain the lock multiple times. The caller must
 // ensure that ReleaseShared is called for each all to AcquireShared.
-func (pl *memLock) AcquireShared(owner uint64) {
-	if owner != uniqueOwner && pl.exclusiveOwner.Load() == owner {
+func (l *memLock) AcquireShared(owner uint64) {
+	if owner != uniqueOwner && l.exclusiveOwner.Load() == owner {
 		// Exclusive lock has already been acquired by this owner.
-		pl.mu.reentrancy++
+		l.mu.reentrancy++
 		return
 	}
 
 	// Block until shared lock is acquired.
-	pl.mu.RLock()
+	l.mu.RLock()
 }
 
 // Release unlocks exclusive write access to the protected resource obtained by
 // a call to Acquire. If the same owner made multiple Acquire calls, the lock
 // isn't released until Release is called the same number of times.
-func (pl *memLock) Release() {
-	if pl.mu.reentrancy > 0 {
-		pl.mu.reentrancy--
+func (l *memLock) Release() {
+	if l.mu.reentrancy > 0 {
+		l.mu.reentrancy--
 		return
 	}
 
 	// No remaining reentrancy, so release lock.
-	pl.exclusiveOwner.Store(uniqueOwner)
-	pl.mu.Unlock()
+	l.exclusiveOwner.Store(uniqueOwner)
+	l.mu.Unlock()
 }
 
 // ReleaseShared unlocks shared read access to the protected resource obtained
 // by a call to AcquireShared. If the same owner made multiple AcquireShared
 // calls, the lock isn't released until ReleaseShared is called the same number
 // of times.
-func (pl *memLock) ReleaseShared() {
-	if pl.mu.reentrancy > 0 {
-		pl.mu.reentrancy--
+func (l *memLock) ReleaseShared() {
+	if l.mu.reentrancy > 0 {
+		l.mu.reentrancy--
 		return
 	}
 
 	// No remaining reentrancy, so release lock.
-	pl.mu.RUnlock()
+	l.mu.RUnlock()
 }

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

@@ -225,7 +225,7 @@ func (tx *memTxn) GetFullVectors(
 			func(vectorWithKey *cspann.VectorWithKey) {
 				// Lock the partition to read its data.
 				memPart := tx.store.lockPartition(
-					treeKey, vectorWithKey.Key.PartitionKey, uniqueOwner, false /* isExclusive */)
+					treeKey, vectorWithKey.Key.PartitionKey, tx.id, false /* isExclusive */)
 				if memPart != nil {
 					defer memPart.lock.ReleaseShared()
 					vectorWithKey.Vector = memPart.lock.partition.Centroid()

pkg/sql/vecindex/cspann/store.go

@@ -158,6 +158,28 @@ type Store interface {
 		expected PartitionMetadata,
 	) (removed bool, err error)
 
+	// TryMoveVector attempts to move a single vector from the source partition to
+	// the target partition, as an atommic operation. If either partition does not
+	// exist, it returns moved=false. If the vector does not exist in the source
+	// partition, or already exists in the target partition, it returns
+	// moved=false.
+	//
+	// Before performing any action, TryMoveVector checks the target partition's
+	// metadata and returns a ConditionFailedError if it is not the same as the
+	// expected metadata.
+	//
+	// NOTE: The source and target partitions must be on the same level. Results
+	// are not defined if this is not true.
+	TryMoveVector(
+		ctx context.Context,
+		treeKey TreeKey,
+		sourcePartitionKey, targetPartitionKey PartitionKey,
+		vec vector.T,
+		childKey ChildKey,
+		valueBytes ValueBytes,
+		expected PartitionMetadata,
+	) (moved 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.