Skip to content

Commit a5bf5b5

Browse files
committed
vecindex: avoid scanning leaf vectors in insert operations
Insert and SearchForInsert operations do not need to scan leaf vectors. Normally, insert operations do not even attempt to load leaf partitions. However, leaf vectors are scanned in the case where the root partition is also a leaf partition (e.g. when there are < ~128 vectors in the K-means tree). This is because, until we read its metadata record, we don't know that the root is a leaf partition. In order to avoid the contention caused by scanning leaf vectors, this commit adds the partition's level into the encoding for each of the KV keys of vectors in that partition: Index Prefix | Prefix Columns | PartitionID | Level | Child PartitionID Encoding the level allows insert operations to skip over leaf vectors by scanning the [prefix/partitionID/2/... - prefix/partitionID/...] span. Epic: CRDB-42943 Release note: None
1 parent 2f36923 commit a5bf5b5

File tree

24 files changed

+447
-260
lines changed

24 files changed

+447
-260
lines changed

docs/generated/metrics/metrics.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1868,7 +1868,7 @@
18681868
<tr><td>APPLICATION</td><td>sql.update.started.count.internal</td><td>Number of SQL UPDATE statements started (internal queries)</td><td>SQL Internal Statements</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
18691869
<tr><td>APPLICATION</td><td>sql.vecindex.fixups_added</td><td>Total number of vector index fixups that have been added to the processor</td><td>Vector Index Fixups Added</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
18701870
<tr><td>APPLICATION</td><td>sql.vecindex.fixups_processed</td><td>Total number of vector index fixups that have been processed</td><td>Vector Index Fixups Processed</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
1871-
<tr><td>APPLICATION</td><td>sql.vecindex.splits</td><td>Total number of vector index partitions split without error</td><td>Vector Index Splits</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
1871+
<tr><td>APPLICATION</td><td>sql.vecindex.successful_splits</td><td>Total number of vector index partitions split without error</td><td>Vector Index Splits</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
18721872
<tr><td>APPLICATION</td><td>sqlliveness.is_alive.cache_hits</td><td>Number of calls to IsAlive that return from the cache</td><td>Calls</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
18731873
<tr><td>APPLICATION</td><td>sqlliveness.is_alive.cache_misses</td><td>Number of calls to IsAlive that do not return from the cache</td><td>Calls</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
18741874
<tr><td>APPLICATION</td><td>sqlliveness.sessions_deleted</td><td>Number of expired sessions which have been deleted</td><td>Sessions</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>

pkg/BUILD.bazel

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -652,6 +652,7 @@ ALL_TESTS = [
652652
"//pkg/sql/vecindex/cspann/quantize:quantize_test",
653653
"//pkg/sql/vecindex/cspann/workspace:workspace_test",
654654
"//pkg/sql/vecindex/cspann:cspann_test",
655+
"//pkg/sql/vecindex/vecencoding:vecencoding_test",
655656
"//pkg/sql/vecindex/vecstore:vecstore_test",
656657
"//pkg/sql/vecindex:vecindex_test",
657658
"//pkg/sql:sql_disallowed_imports_test",
@@ -2360,6 +2361,8 @@ GO_TARGETS = [
23602361
"//pkg/sql/vecindex/cspann/workspace:workspace_test",
23612362
"//pkg/sql/vecindex/cspann:cspann",
23622363
"//pkg/sql/vecindex/cspann:cspann_test",
2364+
"//pkg/sql/vecindex/vecencoding:vecencoding",
2365+
"//pkg/sql/vecindex/vecencoding:vecencoding_test",
23632366
"//pkg/sql/vecindex/vecpb:vecpb",
23642367
"//pkg/sql/vecindex/vecstore:vecstore",
23652368
"//pkg/sql/vecindex/vecstore:vecstore_test",

pkg/cmd/vecbench/sql_provider.go

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,8 +164,36 @@ func (s *SQLProvider) InsertVectors(
164164
func (s *SQLProvider) Search(
165165
ctx context.Context, vec vector.T, maxResults int, beamSize int, stats *cspann.SearchStats,
166166
) ([]cspann.KeyBytes, error) {
167-
// This will be implemented in the next phase
168-
return nil, errors.New("Search not yet implemented for SQL provider")
167+
// Construct a query that searches for similar vectors.
168+
query := fmt.Sprintf(`
169+
SELECT id
170+
FROM %s
171+
ORDER BY embedding <-> $1
172+
LIMIT $2
173+
`, s.tableName)
174+
175+
// Execute the query.
176+
rows, err := s.pool.Query(ctx, query, vec, maxResults)
177+
if err != nil {
178+
return nil, errors.Wrap(err, "executing search query")
179+
}
180+
defer rows.Close()
181+
182+
// Collect the results.
183+
var results []cspann.KeyBytes
184+
for rows.Next() {
185+
var id []byte
186+
if err := rows.Scan(&id); err != nil {
187+
return nil, errors.Wrap(err, "scanning search result")
188+
}
189+
results = append(results, id)
190+
}
191+
192+
if err := rows.Err(); err != nil {
193+
return nil, errors.Wrap(err, "iterating search results")
194+
}
195+
196+
return results, nil
169197
}
170198

171199
// GetMetrics implements the VectorProvider interface.

pkg/sql/opt/exec/execbuilder/testdata/vector_mutation

Lines changed: 29 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -19,65 +19,65 @@ query T kvtrace(redactbytes)
1919
INSERT INTO exec_test VALUES (1, 1, 'foo', '[1, 2, 3]'), (2, 10, 'bar', '[4, 5, 6]'), (3, 15, 'baz', '[7, 8, 9]')
2020
----
2121
CPut /Table/106/1/1/0 -> /TUPLE/2:2:Int/1/1:3:Bytes/foo/
22-
CPut /Table/106/2/1/1/0 -> /BYTES/:redacted:
23-
CPut /Table/106/3/1/1/1/0 -> /BYTES/:redacted:
24-
CPut /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/0 -> /BYTES/:redacted:
22+
CPut /Table/106/2/1/1/1/0 -> /BYTES/:redacted:
23+
CPut /Table/106/3/1/1/1/1/0 -> /BYTES/:redacted:
24+
CPut /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/1/0 -> /BYTES/:redacted:
2525
CPut /Table/106/1/2/0 -> /TUPLE/2:2:Int/10/1:3:Bytes/bar/
26-
CPut /Table/106/2/1/2/0 -> /BYTES/:redacted:
27-
CPut /Table/106/3/10/1/2/0 -> /BYTES/:redacted:
28-
CPut /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/2/0 -> /BYTES/:redacted:
26+
CPut /Table/106/2/1/1/2/0 -> /BYTES/:redacted:
27+
CPut /Table/106/3/10/1/1/2/0 -> /BYTES/:redacted:
28+
CPut /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/1/2/0 -> /BYTES/:redacted:
2929
CPut /Table/106/5/10/2/0 -> /BYTES/
3030
CPut /Table/106/1/3/0 -> /TUPLE/2:2:Int/15/1:3:Bytes/baz/
31-
CPut /Table/106/2/1/3/0 -> /BYTES/:redacted:
32-
CPut /Table/106/3/15/1/3/0 -> /BYTES/:redacted:
33-
CPut /Table/106/4/"\x16\x05\x15\xef\x18\x95\x00\x00\x00 \x00 \x00 "/1/3/0 -> /BYTES/:redacted:
31+
CPut /Table/106/2/1/1/3/0 -> /BYTES/:redacted:
32+
CPut /Table/106/3/15/1/1/3/0 -> /BYTES/:redacted:
33+
CPut /Table/106/4/"\x16\x05\x15\xef\x18\x95\x00\x00\x00 \x00 \x00 "/1/1/3/0 -> /BYTES/:redacted:
3434
CPut /Table/106/5/15/3/0 -> /BYTES/
3535

3636
query T kvtrace(redactbytes)
3737
UPSERT INTO exec_test VALUES (2, 10, 'bar', '[6, 5, 4]'), (4, 20, 'qux', '[10, 11, 12]')
3838
----
3939
Scan /Table/106/1/2/0, /Table/106/1/4/0
4040
Put /Table/106/1/2/0 -> /TUPLE/2:2:Int/10/1:3:Bytes/bar/
41-
Del (locking) /Table/106/2/1/2/0
42-
CPut /Table/106/2/1/2/0 -> /BYTES/:redacted:
43-
Del (locking) /Table/106/3/10/1/2/0
44-
CPut /Table/106/3/10/1/2/0 -> /BYTES/:redacted:
45-
Del (locking) /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/2/0
46-
CPut /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/2/0 -> /BYTES/:redacted:
41+
Del (locking) /Table/106/2/1/1/2/0
42+
CPut /Table/106/2/1/1/2/0 -> /BYTES/:redacted:
43+
Del (locking) /Table/106/3/10/1/1/2/0
44+
CPut /Table/106/3/10/1/1/2/0 -> /BYTES/:redacted:
45+
Del (locking) /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/1/2/0
46+
CPut /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/1/2/0 -> /BYTES/:redacted:
4747
CPut /Table/106/1/4/0 -> /TUPLE/2:2:Int/20/1:3:Bytes/qux/
48-
CPut /Table/106/2/1/4/0 -> /BYTES/:redacted:
49-
CPut /Table/106/3/20/1/4/0 -> /BYTES/:redacted:
50-
CPut /Table/106/4/"\x17\xab\x186\x18{\x00\x00\x00 \x00 \x00 "/1/4/0 -> /BYTES/:redacted:
48+
CPut /Table/106/2/1/1/4/0 -> /BYTES/:redacted:
49+
CPut /Table/106/3/20/1/1/4/0 -> /BYTES/:redacted:
50+
CPut /Table/106/4/"\x17\xab\x186\x18{\x00\x00\x00 \x00 \x00 "/1/1/4/0 -> /BYTES/:redacted:
5151
CPut /Table/106/5/20/4/0 -> /BYTES/
5252

5353
query T kvtrace(redactbytes)
5454
UPDATE exec_test SET encoding = '[3, 2, 1]' WHERE id = 1
5555
----
5656
Scan /Table/106/1/1/0
5757
Put /Table/106/1/1/0 -> /TUPLE/2:2:Int/1/1:3:Bytes/foo/
58-
Del (locking) /Table/106/2/1/1/0
59-
CPut /Table/106/2/1/1/0 -> /BYTES/:redacted:
60-
Del (locking) /Table/106/3/1/1/1/0
61-
CPut /Table/106/3/1/1/1/0 -> /BYTES/:redacted:
62-
Del (locking) /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/0
63-
CPut /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/0 -> /BYTES/:redacted:
58+
Del (locking) /Table/106/2/1/1/1/0
59+
CPut /Table/106/2/1/1/1/0 -> /BYTES/:redacted:
60+
Del (locking) /Table/106/3/1/1/1/1/0
61+
CPut /Table/106/3/1/1/1/1/0 -> /BYTES/:redacted:
62+
Del (locking) /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/1/0
63+
CPut /Table/106/4/"\x16\x84\x17q\x17q\x00\x00\x00 \x00 \x00 "/1/1/1/0 -> /BYTES/:redacted:
6464

6565
query T kvtrace(redactbytes)
6666
UPDATE exec_test SET data = 'fooo' WHERE user_id = 15;
6767
----
6868
Scan /Table/106/5/1{5-6}
6969
Scan /Table/106/1/3/0
7070
Put /Table/106/1/3/0 -> /TUPLE/2:2:Int/15/1:3:Bytes/fooo/
71-
Del (locking) /Table/106/4/"\x16\x05\x15\xef\x18\x95\x00\x00\x00 \x00 \x00 "/1/3/0
72-
CPut /Table/106/4/"\x16\x84\x17q\x17q\x17q\x00\x00\x00 \x00 \x00 \x00 "/1/3/0 -> /BYTES/:redacted:
71+
Del (locking) /Table/106/4/"\x16\x05\x15\xef\x18\x95\x00\x00\x00 \x00 \x00 "/1/1/3/0
72+
CPut /Table/106/4/"\x16\x84\x17q\x17q\x17q\x00\x00\x00 \x00 \x00 \x00 "/1/1/3/0 -> /BYTES/:redacted:
7373

7474
query T kvtrace
7575
DELETE FROM exec_test WHERE user_id = 10
7676
----
7777
Scan /Table/106/5/1{0-1}
7878
Scan /Table/106/1/2/0
79-
Del /Table/106/2/1/2/0
80-
Del /Table/106/3/10/1/2/0
81-
Del /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/2/0
79+
Del /Table/106/2/1/1/2/0
80+
Del /Table/106/3/10/1/1/2/0
81+
Del /Table/106/4/"\x16\x05\x15\xef\x17\xbd\x00\x00\x00 \x00 \x00 "/1/1/2/0
8282
Del /Table/106/5/10/2/0
8383
Del (locking) /Table/106/1/2/0

pkg/sql/rowenc/BUILD.bazel

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ go_library(
3434
"//pkg/sql/sem/tree",
3535
"//pkg/sql/sqlerrors",
3636
"//pkg/sql/types",
37+
"//pkg/sql/vecindex/cspann",
38+
"//pkg/sql/vecindex/vecencoding",
3739
"//pkg/util/buildutil",
3840
"//pkg/util/deduplicate",
3941
"//pkg/util/encoding",
@@ -85,7 +87,7 @@ go_test(
8587
"//pkg/sql/sem/idxtype",
8688
"//pkg/sql/sem/tree",
8789
"//pkg/sql/types",
88-
"//pkg/sql/vecindex/vecstore",
90+
"//pkg/sql/vecindex/vecencoding",
8991
"//pkg/testutils/datapathutils",
9092
"//pkg/testutils/serverutils",
9193
"//pkg/testutils/sqlutils",

pkg/sql/rowenc/index_encoding.go

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ import (
3030
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
3131
"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
3232
"github.com/cockroachdb/cockroach/pkg/sql/types"
33+
"github.com/cockroachdb/cockroach/pkg/sql/vecindex/cspann"
34+
"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecencoding"
3335
"github.com/cockroachdb/cockroach/pkg/util/deduplicate"
3436
"github.com/cockroachdb/cockroach/pkg/util/encoding"
3537
"github.com/cockroachdb/cockroach/pkg/util/intsets"
@@ -1126,6 +1128,9 @@ func encodeTrigramInvertedIndexTableKeys(
11261128
return outKeys, nil
11271129
}
11281130

1131+
// encodeVectorIndexKey creates the key for a vector index leaf partition, minus
1132+
// any suffix primary key columns. See comment at top of
1133+
// sql/vecindex/vecencoding.go for a description of the format.
11291134
func encodeVectorIndexKey(
11301135
index catalog.Index,
11311136
colMap catalog.TableColMap,
@@ -1141,14 +1146,14 @@ func encodeVectorIndexKey(
11411146
// This index is not being updated.
11421147
return nil, nil
11431148
}
1144-
partitionKey := uint64(tree.MustBeDInt(partitionKeyDatum))
1145-
keyPrefix, err = encodeIndexPrefixKeys(index, colMap, values, keyPrefix)
1149+
partitionKey := cspann.PartitionKey(tree.MustBeDInt(partitionKeyDatum))
1150+
key, err = encodeIndexPrefixKeys(index, colMap, values, keyPrefix)
11461151
if err != nil {
11471152
return nil, err
11481153
}
1149-
// TODO(drewk): use the vector index encoding functions (right now, doing so
1150-
// would cause a dependency cycle).
1151-
return encoding.EncodeUvarintAscending(keyPrefix, partitionKey), nil
1154+
key = vecencoding.EncodePartitionKey(key, partitionKey)
1155+
key = vecencoding.EncodePartitionLevel(key, cspann.LeafLevel)
1156+
return key, err
11521157
}
11531158

11541159
// EncodePrimaryIndex constructs the key prefix for the primary index and

pkg/sql/rowenc/index_encoding_test.go

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ import (
3333
"github.com/cockroachdb/cockroach/pkg/sql/sem/idxtype"
3434
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
3535
"github.com/cockroachdb/cockroach/pkg/sql/types"
36-
"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecstore"
36+
"github.com/cockroachdb/cockroach/pkg/sql/vecindex/vecencoding"
3737
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
3838
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
3939
"github.com/cockroachdb/cockroach/pkg/util"
@@ -1292,23 +1292,9 @@ func TestDecodeKeyVals(t *testing.T) {
12921292
}
12931293
}
12941294

1295-
/* Vector indexes are encoded as follows. The index encoder can only write leaf
1296-
keys, so that's what's tested here.
1297-
1298-
Interior (non-leaf) partition KV key:
1299-
┌────────────┬──────────────┬───────────┬─────────────────┐
1300-
│Index Prefix│Prefix Columns│PartitionID│Child PartitionID│
1301-
└────────────┴──────────────┴───────────┴─────────────────┘
1302-
Leaf partition KV key:
1303-
┌────────────┬──────────────┬───────────┬───────────┬────────────────────┐
1304-
│Index Prefix│Prefix Columns│PartitionID│Primary Key│Sentinel Family ID 0│
1305-
└────────────┴──────────────┴───────────┴───────────┴────────────────────┘
1306-
Value:
1307-
┌────────────────────────┬────────────────────────┐
1308-
│Quantized+Encoded Vector│Composite+Stored Columns│
1309-
└────────────────────────┴────────────────────────┘
1310-
*/
1311-
1295+
// See comment at the top of pkg/sql/vecindex/vecstore/encoding.go for more
1296+
// details on the encoding format for vector indexes. The index encoder can only
1297+
// write leaf keys, so that's what's tested here.
13121298
func TestVectorEncoding(t *testing.T) {
13131299
defer leaktest.AfterTest(t)()
13141300

@@ -1330,7 +1316,7 @@ func TestVectorEncoding(t *testing.T) {
13301316
tableDesc := desctestutils.TestingGetPublicTableDescriptor(kvDB, codec, "defaultdb", "prefix_cols")
13311317

13321318
testVector := vector.T{1, 2, 4}
1333-
encodedVector, err := vecstore.EncodeUnquantizedVector([]byte{}, 0, testVector)
1319+
encodedVector, err := vecencoding.EncodeUnquantizedVector([]byte{}, 0, testVector)
13341320
require.NoError(t, err)
13351321

13361322
vh := VectorIndexEncodingHelper{
@@ -1355,8 +1341,11 @@ func TestVectorEncoding(t *testing.T) {
13551341
}
13561342

13571343
expectedKeyMap := map[string][]uint64{
1358-
"simple_idx": {8311, 1, 0},
1359-
"prefix_idx": {3, 2, 8311, 1, 0},
1344+
// Partition key (8311), leaf level (1), suffix key column (1), first
1345+
// byte of primary key (0).
1346+
"simple_idx": {8311, 1, 1, 0},
1347+
// Same as above, prefixed by column c value (3) and column b value (2).
1348+
"prefix_idx": {3, 2, 8311, 1, 1, 0},
13601349
}
13611350

13621351
for _, idx := range tableDesc.PublicNonPrimaryIndexes() {
@@ -1416,7 +1405,7 @@ func TestVectorCompositeEncoding(t *testing.T) {
14161405
tableDesc := desctestutils.TestingGetPublicTableDescriptor(kvDB, codec, "defaultdb", "prefix_cols")
14171406

14181407
testVector := vector.T{1, 2, 4}
1419-
encodedVector, err := vecstore.EncodeUnquantizedVector([]byte{}, 0, testVector)
1408+
encodedVector, err := vecencoding.EncodeUnquantizedVector([]byte{}, 0, testVector)
14201409
require.NoError(t, err)
14211410

14221411
vh := VectorIndexEncodingHelper{
@@ -1475,6 +1464,11 @@ func TestVectorCompositeEncoding(t *testing.T) {
14751464
require.NoError(t, err)
14761465
require.Equal(t, uint64(8311), decodedPartitionID)
14771466

1467+
// Decode the partition Level and ensure that it is correct.
1468+
keyBytes, decodedPartitionLevel, err := encoding.DecodeUvarintAscending(keyBytes)
1469+
require.NoError(t, err)
1470+
require.Equal(t, uint64(1), decodedPartitionLevel)
1471+
14781472
// Decode the primary key and ensure that it is correct.
14791473
pkValue, keyBytes, err := keyside.Decode(&tree.DatumAlloc{}, types.Int, keyBytes, encoding.Ascending)
14801474
require.NoError(t, err)

pkg/sql/vecindex/BUILD.bazel

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ go_test(
5050
"//pkg/security/securitytest",
5151
"//pkg/security/username",
5252
"//pkg/server",
53+
"//pkg/sql",
5354
"//pkg/sql/catalog",
5455
"//pkg/sql/catalog/catenumpb",
5556
"//pkg/sql/catalog/catpb",
@@ -65,6 +66,7 @@ go_test(
6566
"//pkg/sql/vecindex/cspann",
6667
"//pkg/sql/vecindex/cspann/quantize",
6768
"//pkg/sql/vecindex/cspann/testutils",
69+
"//pkg/sql/vecindex/vecencoding",
6870
"//pkg/sql/vecindex/vecpb",
6971
"//pkg/sql/vecindex/vecstore",
7072
"//pkg/testutils/serverutils",

0 commit comments

Comments
 (0)