From 6262246b86a4e07a2fa9afdc29241dafc5a345a3 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Thu, 22 Oct 2026 13:13:38 -0500 Subject: [PATCH] Read a full scan's vectors through one blob handle per chunk A vec0 full scan read each row's vector through vec0_get_vector_data(), which for a flat column looks up the row's chunk, opens the chunk's vectors blob, reads one vector and closes the blob again. A chunk's blob holds chunk_size vectors (3 MiB for 1024 float[768]), and the first read through a newly opened handle walks the blob's overflow pages up to the row's offset, so the opens dominated the scan. The scan already visits rows in (chunk_id, chunk_offset) order, and for a flat column the cursor now keeps one blob handle per vector column, moved to the next chunk with sqlite3_blob_reopen(). Each row's chunk position is looked up when its vector is read: the scan's rowid statement sorts every row when the scan starts, so its positions would miss a delete or insert the same connection runs during the scan. A write that replaces the chunk's row expires the handle, and a rollback of a transaction that changed the schema invalidates it (SQLITE_ABORT, or SQLITE_ABORT_ROLLBACK with extended result codes on); the read is then retried once on a newly opened handle. Any other failure to move the handle, such as an interrupt, fails the read. Rescore, IVF and DiskANN columns keep their vectors in tables of their own and are read row by row as before. The position lookup and that failure handling follow photostructure/sqlite-vec@f793711 by Matthew McEachen, which corrected this change after it merged there as photostructure/sqlite-vec#1. On a 10,000-row float[768] table (medians of 3 interleaved runs of 7 queries), selecting every vector takes 12.6 ms instead of 762.4 ms in a database file and 5.4 ms instead of 104.7 ms in memory, and vec_distance_cosine() against each 33.8 ms instead of 853.5 ms and 24.5 ms instead of 111.8 ms. tests/test-fullscan.py checks every row's vectors against point lookups across chunks with gaps and two vector columns, scans rescore and DiskANN columns beside a flat one, reads a chunk rewritten in the middle of a scan, fails on a row deleted during the scan instead of returning another row's vector, and fails a scan cancelled while its handle moves to the next chunk; tests/test-unit.c scans across a rollback that changed the schema, with extended result codes off and on. --- sqlite-vec.c | 113 ++++++++++++++++++++++++++++++++- tests/test-fullscan.py | 141 +++++++++++++++++++++++++++++++++++++++++ tests/test-unit.c | 53 ++++++++++++++++ 3 files changed, 306 insertions(+), 1 deletion(-) create mode 100644 tests/test-fullscan.py diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..fe6cf7a7 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -4801,8 +4801,16 @@ int vec0_new_chunk(vec0_vtab *p, sqlite3_value ** partitionKeyValues, i64 *chunk } struct vec0_query_fullscan_data { + // Every rowid, in storage order as of the scan's start. sqlite3_stmt *rowids_stmt; i8 done; + // Per flat vector column, a handle on the vectors blob of the chunk the scan + // last read from. A chunk's blob holds chunk_size vectors, and the first read + // through a newly opened handle walks the blob's overflow pages up to the + // row's offset, so a scan keeps it open across the chunk's rows instead of + // opening it once per row. + sqlite3_blob *vectorBlobs[VEC0_MAX_VECTOR_COLUMNS]; + i64 vectorBlobChunkIds[VEC0_MAX_VECTOR_COLUMNS]; }; void vec0_query_fullscan_data_clear( struct vec0_query_fullscan_data *fullscan_data) { @@ -4813,6 +4821,10 @@ void vec0_query_fullscan_data_clear( sqlite3_finalize(fullscan_data->rowids_stmt); fullscan_data->rowids_stmt = NULL; } + for (int i = 0; i < VEC0_MAX_VECTOR_COLUMNS; i++) { + sqlite3_blob_close(fullscan_data->vectorBlobs[i]); + fullscan_data->vectorBlobs[i] = NULL; + } } struct vec0_query_knn_data { @@ -8352,6 +8364,99 @@ static int vec0Eof(sqlite3_vtab_cursor *cur) { return 1; } +/** + * @brief Read the vector of the full scan's current row from a flat column, + * through the handle the scan keeps on the row's chunk. + * + * The scan visits rows in (chunk_id, chunk_offset) order, so the handle moves + * to the next chunk with sqlite3_blob_reopen() once per chunk. A write that + * replaces the chunk's row expires the handle, and sqlite3_blob_read() answers + * SQLITE_ABORT (https://sqlite.org/c3ref/blob_open.html). A rollback of a + * transaction that changed the schema invalidates it, and both + * sqlite3_blob_read() and sqlite3_blob_reopen() answer SQLITE_ABORT_ROLLBACK + * when the connection has extended result codes on. Either way the read is + * retried once on a newly opened handle. + * + * The row's position is looked up here rather than taken from the scan's + * statement: that statement sorts every row when the scan starts, so its + * positions miss a delete, insert or optimize the same connection runs during + * the scan. + * + * @param outVector set to a buffer the caller frees with sqlite3_free() + */ +static int vec0_fullscan_read_vector(vec0_vtab *p, + struct vec0_query_fullscan_data *f, + int vector_idx, void **outVector, + int *outVectorSize) { + i64 rowid = sqlite3_column_int64(f->rowids_stmt, 0); + i64 chunk_id; + i64 chunk_offset; + int rc = vec0_get_chunk_position(p, rowid, NULL, &chunk_id, &chunk_offset); + if (rc == SQLITE_EMPTY) { + vtab_set_error(&p->base, "Could not find a row with rowid %lld", rowid); + return SQLITE_ERROR; + } + if (rc != SQLITE_OK) { + return rc; + } + int size = vector_column_byte_size(p->vector_columns[vector_idx]); + void *buf = sqlite3_malloc(size); + if (!buf) { + return SQLITE_NOMEM; + } + + sqlite3_blob **blob = &f->vectorBlobs[vector_idx]; + rc = SQLITE_ABORT; + for (int attempt = 0; attempt < 2 && (rc & 0xff) == SQLITE_ABORT; attempt++) { + if (*blob && f->vectorBlobChunkIds[vector_idx] != chunk_id) { + rc = sqlite3_blob_reopen(*blob, chunk_id); + if (rc != SQLITE_OK) { + sqlite3_blob_close(*blob); + *blob = NULL; + } + // An invalidated handle is replaced below; any other failure, such as + // SQLITE_INTERRUPT, fails the read. + if (rc != SQLITE_OK && (rc & 0xff) != SQLITE_ABORT) { + sqlite3_free(buf); + vtab_set_error( + &p->base, + "Could not fetch vector data for %lld, opening blob failed", rowid); + return SQLITE_ERROR; + } + } + if (!*blob) { + rc = sqlite3_blob_open(p->db, p->schemaName, + p->shadowVectorChunksNames[vector_idx], "vectors", + chunk_id, 0, blob); + if (rc != SQLITE_OK) { + *blob = NULL; + sqlite3_free(buf); + vtab_set_error( + &p->base, + "Could not fetch vector data for %lld, opening blob failed", rowid); + return SQLITE_ERROR; + } + } + f->vectorBlobChunkIds[vector_idx] = chunk_id; + rc = sqlite3_blob_read(*blob, buf, size, chunk_offset * size); + if ((rc & 0xff) == SQLITE_ABORT) { + sqlite3_blob_close(*blob); + *blob = NULL; + } + } + if (rc != SQLITE_OK) { + sqlite3_free(buf); + vtab_set_error( + &p->base, + "Could not fetch vector data for %lld, reading from blob failed", + rowid); + return SQLITE_ERROR; + } + *outVector = buf; + *outVectorSize = size; + return SQLITE_OK; +} + static int vec0Column_fullscan(vec0_vtab *pVtab, vec0_cursor *pCur, sqlite3_context *context, int i) { if (!pCur->fullscan_data) { @@ -8367,7 +8472,13 @@ static int vec0Column_fullscan(vec0_vtab *pVtab, vec0_cursor *pCur, void *v; int sz; int vector_idx = vec0_column_idx_to_vector_idx(pVtab, i); - int rc = vec0_get_vector_data(pVtab, rowid, vector_idx, &v, &sz); + // Only a flat column keeps its vectors in _vector_chunks; the other index + // types keep them in tables of their own, read per row. + int rc = + pVtab->vector_columns[vector_idx].index_type == VEC0_INDEX_TYPE_FLAT + ? vec0_fullscan_read_vector(pVtab, pCur->fullscan_data, vector_idx, + &v, &sz) + : vec0_get_vector_data(pVtab, rowid, vector_idx, &v, &sz); if (rc != SQLITE_OK) { return rc; } diff --git a/tests/test-fullscan.py b/tests/test-fullscan.py new file mode 100644 index 00000000..460deaed --- /dev/null +++ b/tests/test-fullscan.py @@ -0,0 +1,141 @@ +import sqlite3 +import struct + +import pytest + + +def _f32(values): + return struct.pack(f"{len(values)}f", *values) + + +def test_fullscan_reads_indexed_columns(db): + # Rescore and DiskANN columns keep their vectors outside _vector_chunks, so + # a scan reads them row by row, beside a flat column read by chunk. + db.execute( + "create virtual table v using vec0(" + "flat float[8], " + "r float[8] indexed by rescore(quantizer=bit), " + "d float[8] indexed by diskann(neighbor_quantizer=binary), " + "chunk_size=8)" + ) + for i in range(1, 21): + db.execute( + "insert into v(rowid, flat, r, d) values (?, ?, ?, ?)", + [i, _f32([i] * 8), _f32([-i] * 8), _f32([i, -i] * 4)], + ) + rows = db.execute("select rowid, flat, r, d from v").fetchall() + assert len(rows) == 20 + for rowid, flat, r, d in rows: + assert flat == _f32([rowid] * 8) + assert r == _f32([-rowid] * 8) + assert d == _f32([rowid, -rowid] * 4) + + +def test_fullscan_reads_every_vector_across_chunks(db): + # Two vector columns, so each keeps its own handle, over chunks with gaps. + db.execute( + "create virtual table v using vec0(a float[3], b float[2], chunk_size=8)" + ) + for i in range(1, 51): + db.execute( + "insert into v(rowid, a, b) values (?, ?, ?)", + [i, _f32([i, i + 0.5, -i]), _f32([2 * i, 1 / i])], + ) + db.execute("delete from v where rowid % 7 = 0") + + scanned = {r[0]: (r[1], r[2]) for r in db.execute("select rowid, a, b from v")} + assert len(scanned) == 50 - 7 + for rowid, vectors in scanned.items(): + point = db.execute("select a, b from v where rowid = ?", [rowid]).fetchone() + assert vectors == (point[0], point[1]) + assert vectors[0] == _f32([rowid, rowid + 0.5, -rowid]) + + +def test_fullscan_reads_a_chunk_rewritten_during_the_scan(db): + # Rewriting a chunk's row expires a blob handle open on it; the scan reads + # the rest of the chunk from the new row, not from what it held before. + db.execute("create virtual table v using vec0(a float[2], chunk_size=8)") + for i in range(1, 17): + db.execute("insert into v(rowid, a) values (?, ?)", [i, _f32([i, 0])]) + first = db.execute("select min(chunk_id) from v_chunks").fetchone()[0] + + scan = db.execute("select rowid, a from v") + head = [scan.fetchone() for _ in range(2)] + rewritten = b"".join(_f32([100 + i, 1]) for i in range(1, 9)) + db.execute( + "update v_vector_chunks00 set vectors = ? where rowid = ?", [rewritten, first] + ) + rest = scan.fetchall() + + rows = {r[0]: struct.unpack("2f", r[1]) for r in head + rest} + assert sorted(rows) == list(range(1, 17)) + # Python's cursor may have stepped to the third row before the update, so + # the rows checked are those it can only have read after it. + for rowid in range(4, 9): + assert rows[rowid] == (100 + rowid, 1) + for rowid in range(9, 17): + assert rows[rowid] == (rowid, 0) + + +def _scan_around(db, write): + # Fetches one row of a full scan, runs `write` on the same connection, then + # fetches the rest. + db.execute("create virtual table v using vec0(a float[2], chunk_size=8)") + for i in range(1, 17): + db.execute("insert into v(rowid, a) values (?, ?)", [i, _f32([i, 0])]) + db.execute("delete from v where rowid between 2 and 7") + scan = db.execute("select rowid, a from v") + head = [scan.fetchone()] + write() + return {r[0]: struct.unpack("2f", r[1]) for r in head + scan.fetchall()} + + +def test_fullscan_errors_on_a_row_deleted_during_the_scan(db): + def write(): + db.execute("delete from v where rowid = 12") + db.execute("insert into v(rowid, a) values (99, ?)", [_f32([99, 99])]) + + # Row 99 takes row 12's slot; the scan must not return its vector as row 12. + with pytest.raises(sqlite3.OperationalError, match="rowid 12"): + _scan_around(db, write) + + +def test_fullscan_keeps_an_interrupt_raised_while_moving_to_the_next_chunk(db): + # A progress handler that returns nonzero once cancels the statement. Moving + # the scan's blob handle to the next chunk runs SQLite's VM, so cancelling + # there must fail the scan like cancelling anywhere else. A cancellation + # inside vec0's own statements is reported as a vec0 error rather than + # "interrupted", so the check is that the scan fails. + db.execute("create virtual table v using vec0(a float[2], chunk_size=8)") + for i in range(1, 17): + db.execute("insert into v(rowid, a) values (?, ?)", [i, _f32([i, 0])]) + + def finish_scan(handler): + scan = db.execute("select rowid, a from v") + for _ in range(7): + scan.fetchone() + db.set_progress_handler(handler, 1) + try: + return scan.fetchall() + finally: + db.set_progress_handler(None, 1) + + calls = 0 + + def count(): + nonlocal calls + calls += 1 + return 0 + + finish_scan(count) + total = calls + for n in range(1, total + 1): + calls = 0 + + def cancel_at_n(): + nonlocal calls + calls += 1 + return calls == n + + with pytest.raises(sqlite3.OperationalError): + finish_scan(cancel_at_n) diff --git a/tests/test-unit.c b/tests/test-unit.c index 83cedd5f..2b23a538 100644 --- a/tests/test-unit.c +++ b/tests/test-unit.c @@ -2075,6 +2075,58 @@ void test_diskann_prune_select_max_neighbors_1() { printf(" All diskann_prune_select_max_neighbors_1 tests passed.\n"); } +// Runs a full scan of a 16-row, two-chunk vec0 table and, after `stopAt` rows, +// rolls back a transaction that changed the schema on the same connection. That +// rollback invalidates the scan's blob handle, which then answers +// SQLITE_ABORT_ROLLBACK when the connection has extended result codes on; the +// scan must replace the handle and return every row, as with them off. +static void fullscan_across_schema_rollback(int extendedCodes, int stopAt) { + sqlite3 *db; + sqlite3_stmt *stmt; + int rc = sqlite3_open(":memory:", &db); + assert(rc == SQLITE_OK); + rc = sqlite3_vec_init(db, NULL, NULL); + assert(rc == SQLITE_OK); + sqlite3_extended_result_codes(db, extendedCodes); + rc = sqlite3_exec( + db, "create virtual table v using vec0(a float[2], chunk_size=8)", 0, 0, + 0); + assert(rc == SQLITE_OK); + for (int i = 1; i <= 16; i++) { + char sql[64]; + snprintf(sql, sizeof sql, "insert into v(rowid, a) values (%d, '[%d, 0]')", + i, i); + rc = sqlite3_exec(db, sql, 0, 0, 0); + assert(rc == SQLITE_OK); + } + rc = sqlite3_prepare_v2(db, "select rowid, a from v", -1, &stmt, 0); + assert(rc == SQLITE_OK); + int rows = 0; + while ((rc = sqlite3_step(stmt)) == SQLITE_ROW) { + rows++; + assert(sqlite3_column_int64(stmt, 0) == rows); + assert(((const float *)sqlite3_column_blob(stmt, 1))[0] == rows); + if (rows == stopAt) { + int rrc = + sqlite3_exec(db, "begin; create table u(y); rollback;", 0, 0, 0); + assert(rrc == SQLITE_OK); + } + } + assert(rc == SQLITE_DONE); + assert(rows == 16); + sqlite3_finalize(stmt); + sqlite3_close(db); +} + +void test_vec0_fullscan_schema_rollback() { + printf("Starting %s...\n", __func__); + for (int extendedCodes = 0; extendedCodes <= 1; extendedCodes++) { + // Mid-chunk the read fails; at a chunk's last row the next reopen does. + fullscan_across_schema_rollback(extendedCodes, 3); + fullscan_across_schema_rollback(extendedCodes, 8); + } +} + int main() { printf("Starting unit tests...\n"); #ifdef SQLITE_VEC_ENABLE_AVX @@ -2124,5 +2176,6 @@ int main() { test_diskann_prune_select_single_candidate(); test_diskann_prune_select_all_identical_distances(); test_diskann_prune_select_max_neighbors_1(); + test_vec0_fullscan_schema_rollback(); printf("All unit tests passed.\n"); }