Check a DELETE's chunk for emptiness through its validity handle

sqlite-vec.c

@@ -4700,7 +4700,7 @@ int vec0_new_chunk(vec0_vtab *p, sqlite3_value ** partitionKeyValues, i64 *chunk
   // shadow tables declare "rowid PRIMARY KEY" without the INTEGER type, so
   // the user-defined "rowid" column is NOT an alias for the internal SQLite
   // rowid (_rowid_). When only appending rows these two happen to stay in
-  // sync, but after a chunk is deleted (vec0Update_Delete_DeleteChunkIfEmpty)
+  // sync, but after a chunk is deleted (vec0Update_Delete_DeleteChunk)
   // and a new one is created, the auto-assigned _rowid_ can diverge from the
   // user "rowid" value. Since sqlite3_blob_open() addresses rows by internal
   // _rowid_, we must explicitly set BOTH _rowid_ and "rowid" to the same
@@ -9469,8 +9469,10 @@ int vec0Update_Insert(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
   return rc;
 }
 
+// Clears the row's validity bit, and sets *chunkEmpty to whether that was the
+// chunk's last set bit.
 int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id,
-                                    u64 chunk_offset) {
+                                    u64 chunk_offset, int *chunkEmpty) {
   int rc, brc;
   sqlite3_blob *blobChunksValidity = NULL;
   char unsigned bx;
@@ -9516,6 +9518,32 @@ int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id,
     goto cleanup;
   }
 
+  // Another set bit in this byte keeps the chunk, as it does for most deletes,
+  // so the rest of the bitmap is read only once this byte is zero.
+  *chunkEmpty = 0;
+  if (result == 0) {
+    int validitySize = sqlite3_blob_bytes(blobChunksValidity);
+    unsigned char *validity = sqlite3_malloc(validitySize);
+    if (!validity) {
+      rc = SQLITE_NOMEM;
+      goto cleanup;
+    }
+    rc = sqlite3_blob_read(blobChunksValidity, validity, validitySize, 0);
+    if (rc == SQLITE_OK) {
+      *chunkEmpty = 1;
+      for (int i = 0; i < validitySize; i++) {
+        if (validity[i] != 0) {
+          *chunkEmpty = 0;
+          break;
+        }
+      }
+    } else {
+      vtab_set_error(&p->base, "could not read validity blob for %s.%s.%lld",
+                     p->schemaName, p->shadowChunksName, chunk_id);
+    }
+    sqlite3_free(validity);
+  }
+
 cleanup:
 
   brc = sqlite3_blob_close(blobChunksValidity);
@@ -9618,52 +9646,10 @@ int vec0Update_Delete_ClearVectors(vec0_vtab *p, i64 chunk_id,
   return SQLITE_OK;
 }
 
-int vec0Update_Delete_DeleteChunkIfEmpty(vec0_vtab *p, i64 chunk_id,
-                                          int *deleted) {
-  int rc, brc;
-  sqlite3_blob *blobValidity = NULL;
-  *deleted = 0;
-
-  rc = sqlite3_blob_open(p->db, p->schemaName, p->shadowChunksName, "validity",
-                         chunk_id, 0, &blobValidity);
-  if (rc != SQLITE_OK) {
-    vtab_set_error(&p->base,
-                   "could not open validity blob for chunk %lld", chunk_id);
-    return SQLITE_ERROR;
-  }
-
-  int validitySize = sqlite3_blob_bytes(blobValidity);
-  unsigned char *validityBuf = sqlite3_malloc(validitySize);
-  if (!validityBuf) {
-    sqlite3_blob_close(blobValidity);
-    return SQLITE_NOMEM;
-  }
-
-  rc = sqlite3_blob_read(blobValidity, validityBuf, validitySize, 0);
-  brc = sqlite3_blob_close(blobValidity);
-  if (rc != SQLITE_OK) {
-    sqlite3_free(validityBuf);
-    return rc;
-  }
-  if (brc != SQLITE_OK) {
-    sqlite3_free(validityBuf);
-    return brc;
-  }
-
-  int allZero = 1;
-  for (int i = 0; i < validitySize; i++) {
-    if (validityBuf[i] != 0) {
-      allZero = 0;
-      break;
-    }
-  }
-  sqlite3_free(validityBuf);
-
-  if (!allZero) {
-    return SQLITE_OK;
-  }
-
-  // All validity bits are zero — delete this chunk and its associated data
+// Deletes a chunk whose last row a DELETE cleared, with its rows in every
+// chunk table.
+int vec0Update_Delete_DeleteChunk(vec0_vtab *p, i64 chunk_id) {
+  int rc;
   char *zSql;
   sqlite3_stmt *stmt;
 
@@ -9734,7 +9720,6 @@ int vec0Update_Delete_DeleteChunkIfEmpty(vec0_vtab *p, i64 chunk_id,
     p->stmtLatestChunk = NULL;
   }
 
-  *deleted = 1;
   return SQLITE_OK;
 }
 
@@ -9883,6 +9868,7 @@ int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
   i64 rowid;
   i64 chunk_id = 0;
   i64 chunk_offset = 0;
+  int chunkEmpty = 0;
 
   if (p->pkIsText) {
     rc = vec0_rowid_from_id(p, idValue, &rowid);
@@ -9918,7 +9904,8 @@ int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
     }
 
     // 2. clear validity bit
-    rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset);
+    rc =
+        vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset, &chunkEmpty);
     if (rc != SQLITE_OK) {
       return rc;
     }
@@ -9968,9 +9955,8 @@ int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
     }
 
     // 8. reclaim chunk if fully empty
-    {
-      int chunkDeleted;
-      rc = vec0Update_Delete_DeleteChunkIfEmpty(p, chunk_id, &chunkDeleted);
+    if (chunkEmpty) {
+      rc = vec0Update_Delete_DeleteChunk(p, chunk_id);
       if (rc != SQLITE_OK) {
         return rc;
       }

tests/test-insert-delete.py

@@ -221,6 +221,34 @@ def test_delete_all_rows_deletes_chunk(db):
     assert row[0] == _f32([9.0, 9.0, 9.0, 9.0])
 
 
+def test_delete_keeps_chunk_while_another_validity_byte_has_rows(db):
+    # With chunk_size=16, rows 1-8 fill the first byte of the validity bitmap
+    # and rows 9-16 the second. Clearing either byte leaves the other's rows,
+    # which keep the chunk.
+    def chunks():
+        return db.execute("select count(*) from v_chunks").fetchone()[0]
+
+    def knn():
+        rows = db.execute(
+            "select rowid from v where emb match ? and k = 16", [_f32([0.0])]
+        )
+        return sorted(r[0] for r in rows)
+
+    for first, rest in [("rowid > 8", "rowid <= 8"), ("rowid <= 8", "rowid > 8")]:
+        db.execute("drop table if exists v")
+        db.execute("create virtual table v using vec0(emb float[1], chunk_size=16)")
+        for i in range(1, 17):
+            db.execute("insert into v(rowid, emb) values (?, ?)", [i, _f32([float(i)])])
+
+        db.execute(f"delete from v where {first}")
+        assert chunks() == 1
+        kept = [r[0] for r in db.execute(f"select rowid from v where {rest}")]
+        assert knn() == sorted(kept) and len(kept) == 8
+
+        db.execute(f"delete from v where {rest}")
+        assert chunks() == 0
+
+
 def test_delete_chunk_multiple_chunks(db):
     db.execute("create virtual table v using vec0(emb float[4], chunk_size=8)")