vecstore: encoding cleanup

No functional change here, just some code cleanup.

Epic: None

Author: mw5h

pkg/sql/vecindex/vecencoding/encoding.go

@@ -57,7 +57,7 @@ Vector KV Value:
 func EncodeMetadataKey(
 	indexPrefix []byte, encodedPrefixCols []byte, partitionKey cspann.PartitionKey,
 ) roachpb.Key {
-	capacity := len(indexPrefix) + len(encodedPrefixCols) + EncodedPartitionKeyLen(partitionKey) + 1
+	capacity := len(indexPrefix) + len(encodedPrefixCols) + EncodedPartitionKeyLen(partitionKey) + EncodedPartitionLevelLen(cspann.InvalidLevel)
 	keyBuffer := make([]byte, 0, capacity)
 	keyBuffer = append(keyBuffer, indexPrefix...)
 	keyBuffer = append(keyBuffer, encodedPrefixCols...)
@@ -132,13 +132,16 @@ func DecodeVectorKey(
 		keyBytes = keyBytes[prefixLen:]
 	}
 
-	partitionKey, keyBytes, err := DecodePartitionKey(keyBytes)
+	// Decode the partition key and level.
+	var partitionKey cspann.PartitionKey
+	partitionKey, keyBytes, err = DecodePartitionKey(keyBytes)
 	if err != nil {
 		return vecIndexKey, err
 	}
 	vecIndexKey.PartitionKey = partitionKey
 
-	level, keyBytes, err := DecodePartitionLevel(keyBytes)
+	var level cspann.Level
+	level, keyBytes, err = DecodePartitionLevel(keyBytes)
 	if err != nil {
 		return vecIndexKey, err
 	}

pkg/sql/vecindex/vecencoding/encoding_test.go

@@ -273,7 +273,7 @@ func TestDecodeVectorKey(t *testing.T) {
 	require.Equal(t, expectedSuffix, indexKey.Suffix)
 }
 
-func TestDecodeKeyPrefixColumns(t *testing.T) {
+func TestDecodeKeyWithPrefix(t *testing.T) {
 	var buf []byte
 
 	// Encode a prefix column.
@@ -301,7 +301,7 @@ func TestDecodeKeyPrefixColumns(t *testing.T) {
 	expectedSuffix := []byte("suffixData")
 	buf = append(buf, expectedSuffix...)
 
-	// Extract the key with one prefix column.
+	// Extract the key with two prefix columns.
 	key, err := vecencoding.DecodeVectorKey(buf, 2)
 	require.NoError(t, err)
 
@@ -313,3 +313,23 @@ func TestDecodeKeyPrefixColumns(t *testing.T) {
 	// Verify that the remaining bytes form the suffix.
 	require.Equal(t, expectedSuffix, key.Suffix)
 }
+
+func TestDecodeVectorKeyNoPrefix(t *testing.T) {
+	var buf []byte
+	// Encode a partition key.
+	partitionKey := cspann.PartitionKey(789)
+	buf = vecencoding.EncodePartitionKey(buf, partitionKey)
+	// Encode a partition level.
+	level := cspann.Level(3)
+	buf = vecencoding.EncodePartitionLevel(buf, level)
+	// Add some suffix bytes.
+	expectedSuffix := []byte("vectorSuffix")
+	buf = append(buf, expectedSuffix...)
+
+	indexKey, err := vecencoding.DecodeVectorKey(buf, 0)
+	require.NoError(t, err)
+	require.Empty(t, indexKey.Prefix)
+	require.Equal(t, partitionKey, indexKey.PartitionKey)
+	require.Equal(t, level, indexKey.Level)
+	require.Equal(t, expectedSuffix, indexKey.Suffix)
+}

pkg/sql/vecindex/vecstore/store_txn.go

@@ -274,17 +274,15 @@ func (tx *Txn) SearchPartitions(
 	for i := range toSearch {
 		metadataKey := vecencoding.EncodeMetadataKey(tx.store.prefix, treeKey, toSearch[i].Key)
 		b.Get(metadataKey)
-		var startKey, endKey roachpb.Key
+		var startKey roachpb.Key
 		if toSearch[i].ExcludeLeafVectors {
 			// Skip past vectors at the leaf level.
 			startKey = vecencoding.EncodePrefixVectorKey(metadataKey, cspann.SecondLevel)
-			endKey = vecencoding.EncodeEndVectorKey(metadataKey)
-			b.Scan(startKey, endKey)
 		} else {
 			startKey = vecencoding.EncodeStartVectorKey(metadataKey)
-			endKey = vecencoding.EncodeEndVectorKey(metadataKey)
-			b.Scan(startKey, endKey)
 		}
+		endKey := vecencoding.EncodeEndVectorKey(metadataKey)
+		b.Scan(startKey, endKey)
 
 		if log.ExpensiveLogEnabled(ctx, 2) {
 			log.VEventf(ctx, 2, "Scan %s", roachpb.Span{Key: startKey, EndKey: endKey}.String())