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

sqlite-vec.c

@@ -9549,8 +9549,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;
@@ -9596,6 +9598,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);
@@ -9612,53 +9640,10 @@ int vec0Update_Delete_ClearValidity(vec0_vtab *p, i64 chunk_id,
   return SQLITE_OK;
 }
 
-// Delete the chunk and its associated vector/metadata shadow table rows
-// when every validity bit in the chunk is zero.
-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;
-  }
-
+// Delete a chunk whose last row a DELETE cleared, with its associated
+// vector/metadata shadow table rows.
+int vec0Update_Delete_DeleteChunk(vec0_vtab *p, i64 chunk_id) {
+  int rc;
   char *zSql;
   sqlite3_stmt *stmt;
 
@@ -9717,7 +9702,6 @@ int vec0Update_Delete_DeleteChunkIfEmpty(vec0_vtab *p, i64 chunk_id,
     p->stmtLatestChunk = NULL;
   }
 
-  *deleted = 1;
   return SQLITE_OK;
 }
 
@@ -9980,7 +9964,8 @@ int vec0Update_Delete(sqlite3_vtab *pVTab, sqlite3_value *idValue) {
     return rc;
   }
 
-  rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset);
+  int chunkEmpty;
+  rc = vec0Update_Delete_ClearValidity(p, chunk_id, chunk_offset, &chunkEmpty);
   if (rc != SQLITE_OK) {
     return rc;
   }
@@ -10022,14 +10007,9 @@ 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 (rc != SQLITE_OK) {
-      return rc;
-    }
+  if (chunkEmpty) {
+    return vec0Update_Delete_DeleteChunk(p, chunk_id);
   }
-
   return SQLITE_OK;
 }
 

tests/test-insert-delete.py

@@ -219,6 +219,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)")