Read a full scan's vectors through one blob handle per chunk

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

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)

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