Merge #144712

144712: rowenc: handle encoding of NULL or missing vectors r=mw5h a=mw5h

Previously if the mutationSearch operator failed to find a partition for a vector because it was NULL or missing (e.g. being moved by fixup), rowenc would panic, causing the mutation to fail. This prevented the insertion of NULL vectors and would sometimes cause deletes to erroneously fail.

This patch allows mutationSearch to signal "partition not found, but it's okay" by sending a NULL partition instead of setting the partition value to nil. The 'nil' value now indicates a bad failure that should never happen whereas a NULL partition value means 'skip this row' ala a partial index.

Fixes: #144621
Release note (bug fix): NULL vectors can now be inserted into tables with vector indexes.

Co-authored-by: Matt White <[email protected]>

Author: craig[bot]

pkg/sql/logictest/testdata/logic_test/vector_index

@@ -552,3 +552,25 @@ statement error vector_search_beam_size cannot be less than 1 or greater than 51
 SET vector_search_beam_size=513
 
 subtest end
+
+subtest test_144621
+
+statement ok
+CREATE TABLE test_144621 (a INT PRIMARY KEY, v VECTOR(3))
+
+statement ok
+CREATE VECTOR INDEX vec_idx ON test_144621 (v)
+
+statement ok
+INSERT INTO test_144621 VALUES (1, NULL)
+
+query I
+SELECT a FROM test_144621;
+----
+1
+
+query I
+SELECT a FROM test_144621@vec_idx ORDER BY v <-> '[1, 2, 3]' LIMIT 1;
+----
+
+subtest end

pkg/sql/row/vector_index.go

@@ -82,29 +82,23 @@ func (vm *VectorIndexUpdateHelper) initImpl(
 			if vm.PutPartitionKeys == nil {
 				vm.PutPartitionKeys = make(map[descpb.IndexID]tree.Datum)
 			}
-			if putPartitionKeys[valIdx] != tree.DNull {
-				vm.PutPartitionKeys[idx.GetID()] = putPartitionKeys[valIdx]
-			}
+			vm.PutPartitionKeys[idx.GetID()] = putPartitionKeys[valIdx]
 		}
 
 		// Retrieve the quantized vector, if it exists.
 		if valIdx < len(putQuantizedVecs) {
 			if vm.PutQuantizedVecs == nil {
 				vm.PutQuantizedVecs = make(map[descpb.IndexID]tree.Datum)
 			}
-			if putQuantizedVecs[valIdx] != tree.DNull {
-				vm.PutQuantizedVecs[idx.GetID()] = putQuantizedVecs[valIdx]
-			}
+			vm.PutQuantizedVecs[idx.GetID()] = putQuantizedVecs[valIdx]
 		}
 
 		// Retrieve the partition key value for delete, if it exists.
 		if valIdx < len(delPartitionKeys) {
 			if vm.DelPartitionKeys == nil {
 				vm.DelPartitionKeys = make(map[descpb.IndexID]tree.Datum)
 			}
-			if delPartitionKeys[valIdx] != tree.DNull {
-				vm.DelPartitionKeys[idx.GetID()] = delPartitionKeys[valIdx]
-			}
+			vm.DelPartitionKeys[idx.GetID()] = delPartitionKeys[valIdx]
 		}
 
 		valIdx++

pkg/sql/rowenc/index_encoding.go

@@ -1382,7 +1382,14 @@ func encodeSecondaryIndexKeyWithKeyPrefix(
 		secondaryIndexKey, err = encodeVectorIndexKey(
 			secondaryIndex, colMap, values, keyPrefix, vh,
 		)
-		secondaryKeys = [][]byte{secondaryIndexKey}
+		// This can happen if the search returns a NULL partition, which means the
+		// partition was not found in the DELETE case or the vector being inserted
+		// is NULL.
+		if secondaryIndexKey == nil {
+			secondaryKeys = [][]byte{}
+		} else {
+			secondaryKeys = [][]byte{secondaryIndexKey}
+		}
 	}
 	return secondaryKeys, containsNull, err
 }

pkg/sql/rowexec/vector_search.go

@@ -322,8 +322,13 @@ func (v *vectorMutationSearchProcessor) Next() (rowenc.EncDatumRow, *execinfrapb
 					break
 				}
 				// It is possible for the search not to find the target index entry, in
-				// which case the result is nil.
-				partitionKey = v.searcher.PartitionKey()
+				// which case the result is nil. In this case, we leave the partition key
+				// as NULL to signal to rowenc that we didn't find a partition key, but that
+				// it's okay since it's expected that this can happen due to fixups moving
+				// vector index entries around.
+				if v.searcher.PartitionKey() != nil {
+					partitionKey = v.searcher.PartitionKey()
+				}
 			}
 		}
 		row = append(row, rowenc.DatumToEncDatum(types.Int, partitionKey))

pkg/sql/vecindex/BUILD.bazel

@@ -52,6 +52,7 @@ go_test(
     deps = [
         "//pkg/base",
         "//pkg/keys",
+        "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/security/username",
@@ -65,6 +66,7 @@ go_test(
         "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/desctestutils",
         "//pkg/sql/catalog/tabledesc",
+        "//pkg/sql/rowenc",
         "//pkg/sql/sem/catid",
         "//pkg/sql/sem/idxtype",
         "//pkg/sql/sem/tree",

pkg/sql/vecindex/vecindex_test.go

@@ -16,9 +16,14 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex"
+	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann"
 	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann/testutils"
+	"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecencoding"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
@@ -117,3 +122,81 @@ func buildIndex(
 
 	wait.Wait()
 }
+
+// TestVecindexDeletion tests that rows can be properly deleted from a vector index.
+func TestVecindexDeletion(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	runner := sqlutils.MakeSQLRunner(sqlDB)
+	defer srv.Stopper().Stop(ctx)
+
+	// Enable vector indexes.
+	runner.Exec(t, `SET CLUSTER SETTING feature.vector_index.enabled = true`)
+
+	// Construct the table.
+	runner.Exec(t, "CREATE TABLE t (id INT PRIMARY KEY, v VECTOR(512), VECTOR INDEX (v))")
+
+	// Load a small set of vectors for testing.
+	vectors := testutils.LoadFeatures(t, 10)
+
+	// Insert the vectors.
+	for i := 0; i < vectors.Count; i++ {
+		runner.Exec(t, "INSERT INTO t (id, v) VALUES ($1, $2)", i, vectors.At(i).String())
+	}
+
+	// Verify the vectors were inserted.
+	var count int
+	runner.QueryRow(t, "SELECT count(*) FROM t").Scan(&count)
+	if count != vectors.Count {
+		t.Errorf("expected %d rows, got %d", vectors.Count, count)
+	}
+
+	// Get the table descriptor to find the vector index ID.
+	var tableID uint32
+	runner.QueryRow(t, "SELECT id FROM system.namespace WHERE name = 't'").Scan(&tableID)
+
+	// Get the index ID from crdb_internal.table_indexes
+	var indexID uint32
+	runner.QueryRow(t, "SELECT index_id FROM crdb_internal.table_indexes WHERE descriptor_id = $1 AND index_name = 't_v_idx'", tableID).Scan(&indexID)
+	if indexID == 0 {
+		t.Fatal("vector index not found")
+	}
+
+	// Start a KV transaction to manually delete the vector index keys.
+	db := srv.DB()
+	txn := db.NewTxn(ctx, "delete-vector-index-keys")
+
+	// Delete all vector index keys for the table.
+	codec := srv.ApplicationLayer().Codec()
+	prefix := rowenc.MakeIndexKeyPrefix(codec, descpb.ID(tableID), descpb.IndexID(indexID))
+	prefix = vecencoding.EncodePartitionKey(prefix, cspann.RootKey)
+	prefix = vecencoding.EncodePartitionLevel(prefix, cspann.LeafLevel)
+	key := roachpb.Key(prefix)
+	if _, err := txn.DelRange(ctx, key.Next(), key.PrefixEnd(), false); err != nil {
+		t.Fatal(err)
+	}
+	if err := txn.Commit(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	// Delete rows one at a time and verify the count after each deletion.
+	for i := 0; i < vectors.Count; i++ {
+		runner.Exec(t, "DELETE FROM t WHERE id = $1", i)
+
+		// Verify the count after each deletion
+		runner.QueryRow(t, "SELECT count(*) FROM t").Scan(&count)
+		expectedCount := vectors.Count - (i + 1)
+		if count != expectedCount {
+			t.Errorf("after deleting id %d: expected %d rows, got %d", i, expectedCount, count)
+		}
+	}
+
+	// Final verification that table is empty
+	runner.QueryRow(t, "SELECT count(*) FROM t").Scan(&count)
+	if count != 0 {
+		t.Errorf("expected 0 rows after all deletions, got %d", count)
+	}
+}