Merge #143886

143886: sql: don't encode Table/Index prefix twice for vector search r=mw5h a=DrewKimball

This commit fixes an oversight in how vector search prefix keys are generated; namely, previously we would include the `/Table/Index` key prefix. The vector search library already adds this prefix during a search, so the resulting search keys were incorrect.

Epic: CRDB-42943

Release note: None

Co-authored-by: Drew Kimball <[email protected]>

Author: craig[bot]

pkg/sql/logictest/testdata/logic_test/vector_index

@@ -265,11 +265,53 @@ DROP TABLE alter_test
 statement ok
 CREATE TABLE exec_test (
   a INT PRIMARY KEY,
+  b INT,
   vec1 VECTOR(3),
-  VECTOR INDEX (vec1)
+  VECTOR INDEX idx1 (vec1),
+  VECTOR INDEX idx2 (b, vec1)
 )
 
-# TODO(drewk): write these tests once execution is supported.
+statement ok
+INSERT INTO exec_test (a, b, vec1) VALUES
+  (1, 1, '[1, 2, 3]'),
+  (2, 1, '[4, 5, 6]'),
+  (3, 2, '[7, 8, 9]'),
+  (4, 2, '[10, 11, 12]'),
+  (5, 2, '[13, 14, 15]'),
+  (6, NULL, '[16, 17, 18]'),
+  (7, NULL, '[1, 1, 1]');
+
+# TODO(143209): write a full set of tests once we can make them deterministic.
+# For now, we can write tests that return every vector with a given prefix.
+query I rowsort
+SELECT a FROM exec_test@idx1 ORDER BY vec1 <-> '[1, 1, 2]' LIMIT 7;
+----
+7
+1
+2
+3
+4
+5
+6
+
+query I rowsort
+SELECT a FROM exec_test@idx2 WHERE b = 1 ORDER BY vec1 <-> '[1, 1, 2]' LIMIT 2;
+----
+1
+2
+
+query I rowsort
+SELECT a FROM exec_test@idx2 WHERE b = 2 ORDER BY vec1 <-> '[1, 1, 2]' LIMIT 3;
+----
+3
+4
+5
+
+query I rowsort
+SELECT a FROM exec_test WHERE b IS NULL ORDER BY vec1 <-> '[1, 1, 2]' LIMIT 3;
+----
+7
+6
 
 statement ok
 DROP TABLE exec_test

pkg/sql/span/BUILD.bazel

@@ -41,6 +41,7 @@ go_test(
     embed = [":span"],
     deps = [
         "//pkg/base",
+        "//pkg/keys",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
         "//pkg/server",
@@ -51,6 +52,7 @@ go_test(
         "//pkg/sql/catalog/fetchpb",
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/opt/constraint",
+        "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/testutils/serverutils",
         "//pkg/util/encoding",

pkg/sql/span/span_builder.go

@@ -299,7 +299,7 @@ func (s *Builder) appendSpansFromConstraintSpan(
 	var err error
 	var containsNull bool
 	// Encode each logical part of the start key.
-	span.Key, containsNull, err = s.encodeConstraintKey(cs.StartKey())
+	span.Key, containsNull, err = s.encodeConstraintKey(cs.StartKey(), true /* includePrefix */)
 	if err != nil {
 		return nil, err
 	}
@@ -312,7 +312,7 @@ func (s *Builder) appendSpansFromConstraintSpan(
 		span.Key = span.Key.PrefixEnd()
 	}
 	// Encode each logical part of the end key.
-	span.EndKey, _, err = s.encodeConstraintKey(cs.EndKey())
+	span.EndKey, _, err = s.encodeConstraintKey(cs.EndKey(), true /* includePrefix */)
 	if err != nil {
 		return nil, err
 	}
@@ -337,13 +337,18 @@ func (s *Builder) appendSpansFromConstraintSpan(
 
 // encodeConstraintKey encodes each logical part of a constraint.Key into a
 // roachpb.Key.
+//
+// includePrefix is true if the KeyPrefix bytes should be included in the
+// returned key.
 func (s *Builder) encodeConstraintKey(
-	ck constraint.Key,
+	ck constraint.Key, includePrefix bool,
 ) (key roachpb.Key, containsNull bool, _ error) {
 	if ck.IsEmpty() {
 		return key, containsNull, nil
 	}
-	key = append(key, s.KeyPrefix...)
+	if includePrefix {
+		key = append(key, s.KeyPrefix...)
+	}
 	for i := 0; i < ck.Length(); i++ {
 		val := ck.Value(i)
 		if val == tree.DNull {
@@ -494,8 +499,10 @@ func (s *Builder) KeysFromVectorPrefixConstraint(
 		if !span.HasSingleKey(ctx, s.evalCtx) {
 			return nil, errors.AssertionFailedf("constraint span %s does not have a single key", span)
 		}
+		// Do not include the /Table/Index prefix bytes - we only want the portion
+		// of the prefix that corresponds to the prefix columns.
 		var err error
-		prefixKeys[i], _, err = s.encodeConstraintKey(span.StartKey())
+		prefixKeys[i], _, err = s.encodeConstraintKey(span.StartKey(), false /* includePrefix */)
 		if err != nil {
 			return nil, err
 		}

pkg/sql/span/span_builder_test.go

@@ -10,21 +10,29 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/fetchpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
+	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/encoding"
 	"github.com/stretchr/testify/require"
 )
 
 func TestBuilder_EncodeConstraintKey(t *testing.T) {
+	const (
+		tableID descpb.ID      = 100
+		indexID descpb.IndexID = 2
+	)
 	var (
-		colDirs1  = []string{"asc", "asc", "asc"}
-		colDirs2  = []string{"desc", "asc", "desc"}
-		intDatum1 = tree.NewDInt(1)
-		intDatum2 = tree.NewDInt(2)
-		textDatum = tree.NewDString("foo")
+		tableIndexBytes = rowenc.MakeIndexKeyPrefix(keys.SystemSQLCodec, tableID, indexID)
+		colDirs1        = []string{"asc", "asc", "asc"}
+		colDirs2        = []string{"desc", "asc", "desc"}
+		intDatum1       = tree.NewDInt(1)
+		intDatum2       = tree.NewDInt(2)
+		textDatum       = tree.NewDString("foo")
 	)
 	for tcIdx, tc := range []struct {
 		dirs []string
@@ -63,22 +71,36 @@ func TestBuilder_EncodeConstraintKey(t *testing.T) {
 		},
 	} {
 		t.Run(fmt.Sprintf("case %d", tcIdx+1), func(t *testing.T) {
-			b := Builder{}
-			b.keyAndPrefixCols = make([]fetchpb.IndexFetchSpec_KeyColumn, len(tc.dirs))
-			valDirs := make([]encoding.Direction, len(tc.dirs))
-			for i, dir := range tc.dirs {
-				if dir == "asc" {
-					b.keyAndPrefixCols[i].Direction = catenumpb.IndexColumn_ASC
-					valDirs[i] = encoding.Ascending
-				} else {
-					b.keyAndPrefixCols[i].Direction = catenumpb.IndexColumn_DESC
-					valDirs[i] = encoding.Descending
-				}
+			for _, usePrefix := range []bool{true, false} {
+				t.Run(fmt.Sprintf("usePrefix=%t", usePrefix), func(t *testing.T) {
+					b := Builder{
+						KeyPrefix:        tableIndexBytes,
+						keyAndPrefixCols: make([]fetchpb.IndexFetchSpec_KeyColumn, len(tc.dirs)),
+					}
+					valDirs := make([]encoding.Direction, len(tc.dirs))
+					for i, dir := range tc.dirs {
+						if dir == "asc" {
+							b.keyAndPrefixCols[i].Direction = catenumpb.IndexColumn_ASC
+							valDirs[i] = encoding.Ascending
+						} else {
+							b.keyAndPrefixCols[i].Direction = catenumpb.IndexColumn_DESC
+							valDirs[i] = encoding.Descending
+						}
+					}
+					if usePrefix {
+						prefixDirs := []encoding.Direction{encoding.Ascending, encoding.Ascending}
+						valDirs = append(prefixDirs, valDirs...)
+					}
+					outKey, _, err := b.encodeConstraintKey(tc.in, usePrefix)
+					require.NoError(t, err)
+					vals, _ := encoding.PrettyPrintValuesWithTypes(valDirs, outKey)
+					expected := tc.out
+					if usePrefix && !tc.in.IsEmpty() {
+						expected = fmt.Sprintf("/%d/%d%s", tableID, indexID, expected)
+					}
+					require.Equal(t, expected, "/"+strings.Join(vals, "/"))
+				})
 			}
-			outKey, _, err := b.encodeConstraintKey(tc.in)
-			require.NoError(t, err)
-			vals, _ := encoding.PrettyPrintValuesWithTypes(valDirs, outKey)
-			require.Equal(t, tc.out, "/"+strings.Join(vals, "/"))
 		})
 	}
 }