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

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Changed
+
+- Reading a vector column without a KNN query, as in `select embedding from v`, reads each chunk's vectors through one blob handle instead of opening the chunk's blob again for every row: selecting every vector of a 10,000-row `float[768]` table took 11 ms instead of 844 ms, and `select vec_distance_cosine(embedding, ?) from v` 30 ms instead of 807 ms.
+
 ### Fixed
 
 - A KNN query that takes its query vector, `k`, or `mmr_lambda` from another table of a join now uses it. `select v.rowid from p join v on v.embedding match ? and v.k = p.k` failed with `A LIMIT or 'k = ?' constraint is required`; the same join with `v.k = 3 and v.mmr_lambda = p.lambda` returned no rows; and `select v.rowid from v join p on p.id = v.tag where v.embedding match p.vector and v.k = 2`, with `p.id` an integer primary key, could return no rows. These joins now run one KNN per row of the other table, so a filter on the `vec0` table whose value comes from that table applies before the k nearest rows are chosen. Ported in part from the unmerged [vlasky#14](https://github.com/vlasky/sqlite-vec/pull/14) ([asg017#331](https://github.com/asg017/sqlite-vec/pull/331)), which also applies filters whose value comes from a joined table before the k nearest rows are chosen; that part is not ported, and ARCHITECTURE.md's idxStr section says why.

sqlite-vec.c

@@ -5020,8 +5020,15 @@ int vec0_new_chunk(vec0_vtab *p, sqlite3_value **partitionKeyValues,
 }
 
 struct vec0_query_fullscan_data {
+  // Every row as (rowid, chunk_id, chunk_offset), in storage order.
   sqlite3_stmt *rowids_stmt;
   i8 done;
+  // Per 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 opening it seeks
+  // through its overflow pages, 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) {
@@ -5032,6 +5039,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 {
@@ -9015,7 +9026,7 @@ int vec0Filter_fullscan(vec0_vtab *p, vec0_cursor *pCur) {
   }
   memset(fullscan_data, 0, sizeof(*fullscan_data));
 
-  zSql = sqlite3_mprintf(" SELECT rowid "
+  zSql = sqlite3_mprintf(" SELECT rowid, chunk_id, chunk_offset "
                          " FROM " VEC0_SHADOW_ROWIDS_NAME
                          " ORDER by chunk_id, chunk_offset ",
                          p->schemaName, p->tableName);
@@ -9226,6 +9237,72 @@ static int vec0Eof(sqlite3_vtab_cursor *cur) {
   return 1;
 }
 
+/**
+ * @brief Read the vector of the full scan's current row, 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 (sqlite3_blob_read() answers
+ * SQLITE_ABORT, https://sqlite.org/c3ref/blob_open.html), and the read is
+ * retried once on a newly opened handle.
+ *
+ * @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 = sqlite3_column_int64(f->rowids_stmt, 1);
+  i64 chunk_offset = sqlite3_column_int64(f->rowids_stmt, 2);
+  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];
+  int rc = SQLITE_ABORT;
+  for (int attempt = 0; attempt < 2 && rc == SQLITE_ABORT; attempt++) {
+    if (*blob && f->vectorBlobChunkIds[vector_idx] != chunk_id &&
+        sqlite3_blob_reopen(*blob, chunk_id) != SQLITE_OK) {
+      sqlite3_blob_close(*blob);
+      *blob = NULL;
+    }
+    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 vec0_error_rc(rc);
+      }
+    }
+    f->vectorBlobChunkIds[vector_idx] = chunk_id;
+    rc = sqlite3_blob_read(*blob, buf, size, chunk_offset * size);
+    if (rc == 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) {
@@ -9244,7 +9321,8 @@ 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);
+    int rc = vec0_fullscan_read_vector(pVtab, pCur->fullscan_data, vector_idx,
+                                       &v, &sz);
     if (rc != SQLITE_OK) {
       return rc;
     }

tests/test-fullscan.py

@@ -0,0 +1,51 @@
+import struct
+
+
+def _f32(values):
+    return struct.pack(f"{len(values)}f", *values)
+
+
+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)