Merge #141109 #141143

141109: sql/vecindex: add SearchForDelete method to VectorIndex r=drewkimball a=andy-kimball

Previously, the Delete method had both searched for a vector's partition and handle deleting the vector from that partition. However, the execution engine wants to find the partition without actually deleting the vector, since it will do the deletion as part of its own KV batching logic.

The new SearchForDelete method handles just the search portion of deletion, returning the SearchResult containing the partition key. In addition, the Delete method has been refactored to use SearchForDelete internally.

Epic: CRDB-42943

Release note: None

141143: sql: skip TestBackfillWithProtectedTS r=fqazi a=msbutler

I've observed it flake bors several times.

Informs #139862

Release note: none

Co-authored-by: Andrew Kimball <[email protected]>
Co-authored-by: Michael Butler <[email protected]>

Author: 3 people

pkg/sql/backfill_protected_timestamp_test.go

@@ -233,6 +233,7 @@ func TestBackfillWithProtectedTS(t *testing.T) {
 	skip.UnderDeadlock(t, "test takes too long")
 	skip.UnderStress(t, "test takes too long")
 	skip.UnderRace(t, "test takes too long")
+	skip.WithIssue(t, 139862)
 
 	ctx := context.Background()
 	backfillQueryWait := make(chan struct{})

pkg/sql/vecindex/testdata/search-for-delete.ddt

@@ -0,0 +1,155 @@
+# ----------
+# Search tree with multiple partitions and duplicate data.
+# ----------
+new-index dims=2 min-partition-size=1 max-partition-size=4 beam-size=2
+vec1: (1, 2)
+vec2: (7, 4)
+vec3: (4, 3)
+vec4: (2, 8)
+vec5: (5, 1)
+vec6: (3, 6)
+vec7: (6, 5)
+vec8: (8, 3)
+vec9: (1, 7)
+vec10: (1, 7)
+vec11: (1, 7)
+vec12: (1, 7)
+vec13: (1, 7)
+----
+• 1 (3.2083, 5.125)
+│
+├───• 10 (1.5, 7)
+│   │
+│   ├───• 9 (1, 7)
+│   │   │
+│   │   ├───• vec9 (1, 7)
+│   │   └───• vec13 (1, 7)
+│   │
+│   ├───• 8 (1, 7)
+│   │   │
+│   │   ├───• vec12 (1, 7)
+│   │   ├───• vec10 (1, 7)
+│   │   └───• vec11 (1, 7)
+│   │
+│   └───• 7 (2.5, 7)
+│       │
+│       ├───• vec4 (2, 8)
+│       └───• vec6 (3, 6)
+│
+└───• 11 (4.9167, 3.25)
+    │
+    ├───• 5 (6.5, 4.5)
+    │   │
+    │   ├───• vec7 (6, 5)
+    │   ├───• vec2 (7, 4)
+    │   └───• vec8 (8, 3)
+    │
+    └───• 3 (3.3333, 2)
+        │
+        ├───• vec3 (4, 3)
+        ├───• vec5 (5, 1)
+        └───• vec1 (1, 2)
+
+# Search for vector that exists.
+search-for-delete
+vec1
+----
+vec1: partition 3
+
+# Search for vector that does not exist in tree.
+search-for-delete
+vec100: (10, 10)
+----
+vec100: vector not found
+
+# Search for vector that exists, but without a matching key.
+search-for-delete
+vec100: (1, 2)
+----
+vec100: vector not found
+
+# Search for duplicate vector.
+search-for-delete
+vec12
+----
+vec12: partition 8
+
+# Delete vector from store, but leave it in index.
+delete not-found
+vec1
+----
+• 1 (3.2083, 5.125)
+│
+├───• 10 (1.5, 7)
+│   │
+│   ├───• 9 (1, 7)
+│   │   │
+│   │   ├───• vec9 (1, 7)
+│   │   └───• vec13 (1, 7)
+│   │
+│   ├───• 8 (1, 7)
+│   │   │
+│   │   ├───• vec12 (1, 7)
+│   │   ├───• vec10 (1, 7)
+│   │   └───• vec11 (1, 7)
+│   │
+│   └───• 7 (2.5, 7)
+│       │
+│       ├───• vec4 (2, 8)
+│       └───• vec6 (3, 6)
+│
+└───• 11 (4.9167, 3.25)
+    │
+    ├───• 5 (6.5, 4.5)
+    │   │
+    │   ├───• vec7 (6, 5)
+    │   ├───• vec2 (7, 4)
+    │   └───• vec8 (8, 3)
+    │
+    └───• 3 (3.3333, 2)
+        │
+        ├───• vec3 (4, 3)
+        ├───• vec5 (5, 1)
+        └───• vec1 (MISSING)
+
+# Try to find the missing vector.
+search-for-delete
+vec1: (1, 2)
+----
+vec1: vector not found
+
+# Vector should now be gone from the index.
+format-tree
+----
+• 1 (3.2083, 5.125)
+│
+├───• 10 (1.5, 7)
+│   │
+│   ├───• 9 (1, 7)
+│   │   │
+│   │   ├───• vec9 (1, 7)
+│   │   └───• vec13 (1, 7)
+│   │
+│   ├───• 8 (1, 7)
+│   │   │
+│   │   ├───• vec12 (1, 7)
+│   │   ├───• vec10 (1, 7)
+│   │   └───• vec11 (1, 7)
+│   │
+│   └───• 7 (2.5, 7)
+│       │
+│       ├───• vec4 (2, 8)
+│       └───• vec6 (3, 6)
+│
+└───• 11 (4.9167, 3.25)
+    │
+    ├───• 5 (6.5, 4.5)
+    │   │
+    │   ├───• vec7 (6, 5)
+    │   ├───• vec2 (7, 4)
+    │   └───• vec8 (8, 3)
+    │
+    └───• 3 (3.3333, 2)
+        │
+        ├───• vec3 (4, 3)
+        └───• vec5 (5, 1)

pkg/sql/vecindex/vector_index.go

@@ -120,8 +120,8 @@ type searchContext struct {
 	tempVectorsWithKeys []vecstore.VectorWithKey
 }
 
-// VectorIndex implements the C-SPANN algorithm, which adapts Microsoft’s SPANN
-// and SPFresh algorithms to work well with CockroachDB’s unique distributed
+// VectorIndex implements the C-SPANN algorithm, which adapts Microsoft's SPANN
+// and SPFresh algorithms to work well with CockroachDB's unique distributed
 // architecture. This enables CockroachDB to efficiently answer approximate
 // nearest neighbor (ANN) queries with high accuracy, low latency, and fresh
 // results, with millions or even billions of indexed vectors. In a departure
@@ -314,60 +314,17 @@ func (vi *VectorIndex) Insert(
 func (vi *VectorIndex) Delete(
 	ctx context.Context, txn vecstore.Txn, vector vector.T, key vecstore.PrimaryKey,
 ) error {
-	// Potentially throttle delete operation if background work is falling behind.
-	if err := vi.fixups.DelayInsertOrDelete(ctx); err != nil {
+	result, err := vi.SearchForDelete(ctx, txn, vector, key)
+	if err != nil {
 		return err
 	}
-
-	// Search for the vector in the index.
-	searchCtx := searchContext{
-		Txn:      txn,
-		Original: vector,
-		Level:    vecstore.LeafLevel,
-		Options: SearchOptions{
-			SkipRerank:  vi.options.DisableErrorBounds,
-			UpdateStats: true,
-		},
+	if result == nil {
+		return nil
 	}
-	searchCtx.Ctx = internal.WithWorkspace(ctx, &searchCtx.Workspace)
-
-	// Randomize the vector.
-	tempRandomized := searchCtx.Workspace.AllocVector(vi.quantizer.GetDims())
-	defer searchCtx.Workspace.FreeVector(tempRandomized)
-	searchCtx.Randomized = vi.randomizeVector(vector, tempRandomized)
-
-	searchSet := vecstore.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 {
-		searchCtx.Options.BaseBeamSize = baseBeamSize
-
-		err := vi.searchHelper(&searchCtx, &searchSet)
-		if err != nil {
-			return err
-		}
-		results := searchSet.PopUnsortedResults()
-		if len(results) == 0 {
-			// Retry search with significantly higher beam size.
-			if baseBeamSize == vi.options.BaseBeamSize {
-				baseBeamSize *= 8
-				continue
-			}
-			return nil
-		}
-
-		// Remove the vector from its partition in the store.
-		_, err = vi.removeFromPartition(ctx, txn, results[0].ParentPartitionKey, results[0].ChildKey)
-		if errors.Is(err, vecstore.ErrRestartOperation) {
-			// If store requested the operation be retried, then re-run the
-			// search and delete.
-			continue
-		}
-		return err
-	}
+	// Remove the vector from its partition in the store.
+	_, err = vi.removeFromPartition(ctx, txn, result.ParentPartitionKey, result.ChildKey)
+	return err
 }
 
 // Search finds vectors in the index that are closest to the given query vector
@@ -438,6 +395,59 @@ func (vi *VectorIndex) SearchForInsert(
 	return result, nil
 }
 
+// SearchForDelete finds the leaf partition containing the vector to be deleted.
+// It returns a single search result containing the key of that partition, or
+// nil if the vector cannot be found. This is useful for callers that directly
+// delete KV rows rather than using this library to do it.
+func (vi *VectorIndex) SearchForDelete(
+	ctx context.Context, txn vecstore.Txn, vector vector.T, key vecstore.PrimaryKey,
+) (*vecstore.SearchResult, error) {
+	// Potentially throttle operation if background work is falling behind.
+	if err := vi.fixups.DelayInsertOrDelete(ctx); err != nil {
+		return nil, err
+	}
+
+	searchCtx := searchContext{
+		Txn:      txn,
+		Original: vector,
+		Level:    vecstore.LeafLevel,
+		Options: SearchOptions{
+			SkipRerank:  vi.options.DisableErrorBounds,
+			UpdateStats: true,
+		},
+	}
+	searchCtx.Ctx = internal.WithWorkspace(ctx, &searchCtx.Workspace)
+
+	// Randomize the vector.
+	tempRandomized := searchCtx.Workspace.AllocVector(vi.quantizer.GetDims())
+	defer searchCtx.Workspace.FreeVector(tempRandomized)
+	searchCtx.Randomized = vi.randomizeVector(vector, tempRandomized)
+
+	searchCtx.tempSearchSet = vecstore.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++ {
+		searchCtx.Options.BaseBeamSize = baseBeamSize
+
+		err := vi.searchHelper(&searchCtx, &searchCtx.tempSearchSet)
+		if err != nil {
+			return nil, err
+		}
+		results := searchCtx.tempSearchSet.PopUnsortedResults()
+		if len(results) == 0 {
+			// Retry search with significantly higher beam size.
+			baseBeamSize *= 8
+		} else {
+			return &results[0], nil
+		}
+	}
+
+	return nil, nil
+}
+
 // SuspendFixups suspends background fixup processing until ProcessFixups is
 // explicitly called. It is used for testing.
 func (vi *VectorIndex) SuspendFixups() {