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"); }