Use tantivy's AliveBitSet instead of our DeletedBitset for storing deleted tantivy ids (#25520)

Using tantivy's `AliveBitSet` instead of the `DeletedBitset` we use for vector enables us to reuse their segment merge logic.

Also splits `MemoryIdAndDeletionTracker` into `SearchMemoryIdTracker`, which is in charge of building the id tracker, and `MemoryDeletionTracker`, which tracks deletes in other segments.

GitOrigin-RevId: ec770f176d56c847ce4968fad01ddf1cd72e0188

Author: emmaling27

Cargo.lock

@@ -109,8 +109,8 @@ sourcemap = "7"
 strum = { version = "0.26", features = [ "derive" ] }
 sucds = { version = "0.8.1", features = [ "intrinsics" ] }
 syn = { version = "2.0", features = [ "full" ] }
-tantivy = { git = "https://github.com/get-convex/tantivy", rev = "a0f4d3c3843dd255ac6a7210830b3f91cf16db8f" }
-tantivy-common = { git = "https://github.com/get-convex/tantivy", rev = "a0f4d3c3843dd255ac6a7210830b3f91cf16db8f" }
+tantivy = { git = "https://github.com/get-convex/tantivy", rev = "b0bc7a40d54c92c8446098f20082756f218cdd9d" }
+tantivy-common = { git = "https://github.com/get-convex/tantivy", rev = "b0bc7a40d54c92c8446098f20082756f218cdd9d" }
 tempfile = "3"
 thiserror = "1"
 thousands = "0.2.0"

Cargo.toml

@@ -16,8 +16,6 @@ use byteorder::{
 };
 use csf::ls::Map as CsfMap;
 
-use crate::deleted_bitset::DeletedBitset;
-
 /// Version 1 of the id table has the following format:
 /// ```
 /// [ version ] [ count ] [ index_len ] [ ID ]* [ index ]
@@ -42,35 +40,23 @@ pub struct StaticIdTracker {
     /// Convex IDs in search/vector index id order.
     id_buf: Vec<u8>,
     csf_map: CsfMap,
-    deleted: DeletedBitset,
 }
 
 impl StaticIdTracker {
-    pub fn load_from_path(
-        id_table_path: PathBuf,
-        deleted_bitset: DeletedBitset,
-    ) -> anyhow::Result<Self> {
+    pub fn load_from_path(id_table_path: PathBuf) -> anyhow::Result<Self> {
         let uuid_file = File::open(id_table_path)?;
-        StaticIdTracker::load(
-            (
-                uuid_file.metadata()?.len() as usize,
-                BufReader::new(uuid_file),
-            ),
-            deleted_bitset,
-        )
+        StaticIdTracker::load((
+            uuid_file.metadata()?.len() as usize,
+            BufReader::new(uuid_file),
+        ))
     }
 
-    pub fn load(
-        uuid_file: (usize, impl Read),
-        deleted_bitset: DeletedBitset,
-    ) -> anyhow::Result<Self> {
+    pub fn load(uuid_file: (usize, impl Read)) -> anyhow::Result<Self> {
         let (count, uuid_buf, csf_map) = Self::load_ids(uuid_file.0, uuid_file.1)?;
-        anyhow::ensure!(count == deleted_bitset.len());
         Ok(Self {
             count,
             id_buf: uuid_buf,
             csf_map,
-            deleted: deleted_bitset,
         })
     }
 
@@ -120,10 +106,6 @@ impl StaticIdTracker {
         self.count
     }
 
-    pub fn deleted(&self) -> &DeletedBitset {
-        &self.deleted
-    }
-
     pub fn lookup(&self, convex_id: [u8; 16]) -> Option<u32> {
         let index_id = self.csf_map.get(&convex_id);
         let found_convex_id = self.get_convex_id(index_id as usize)?;

crates/common/src/id_tracker.rs

@@ -3,8 +3,8 @@ use std::{
     fmt,
 };
 
-use common::deleted_bitset::DeletedBitset;
 use tantivy::{
+    fastfield::AliveBitSet,
     query::{
         intersect_scorers,
         BooleanQuery,
@@ -151,26 +151,31 @@ impl Weight for ConvexSearchWeight {
 #[derive(Clone)]
 pub struct DeletedDocuments {
     pub memory_deleted: BTreeSet<DocId>,
-    pub segment_deleted: DeletedBitset,
+    pub segment_alive_bitset: AliveBitSet,
 }
 
 impl fmt::Debug for DeletedDocuments {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("DeletedDocuments")
             .field("memory_deleted", &self.memory_deleted)
             .field("segment_deleted", &"<bitset>")
-            .field("num_segment_deleted", &self.segment_deleted.num_deleted())
+            .field(
+                "num_segment_alive",
+                &self.segment_alive_bitset.num_alive_docs(),
+            )
             .finish()
     }
 }
 
 impl DeletedDocuments {
     pub fn contains(&self, doc: DocId) -> bool {
-        self.memory_deleted.contains(&doc) || self.segment_deleted.is_deleted(doc)
+        self.memory_deleted.contains(&doc) || self.segment_alive_bitset.is_deleted(doc)
     }
 
-    pub fn len(&self) -> usize {
-        self.memory_deleted.len() + self.segment_deleted.num_deleted()
+    pub fn approximate_num_alive_docs(&self) -> usize {
+        self.segment_alive_bitset
+            .num_alive_docs()
+            .saturating_sub(self.memory_deleted.len())
     }
 }
 
@@ -225,9 +230,7 @@ impl<T: DocSet> DocSet for ExcludeDeleted<T> {
     }
 
     fn size_hint(&self) -> u32 {
-        self.docset
-            .size_hint()
-            .saturating_sub(self.deleted_documents.len() as u32)
+        self.deleted_documents.approximate_num_alive_docs() as u32
     }
 }
 

crates/search/src/convex_query.rs

@@ -292,18 +292,18 @@ impl MutableFragmentedSegmentMetadata {
         download_single_file_zip(&original.deleted_bitset_key, &deleted_bitset_path, storage)
             .await?;
 
-        let deleted_bitset = DeletedBitset::load_from_path(deleted_bitset_path)?;
+        let deleted = DeletedBitset::load_from_path(deleted_bitset_path)?;
 
         // Clone is a bit of a hack here because these two deleted bitsets may become
         // inconsistent if one or more vectors are deleted via maybe_delete.
         // For now we don't care about the inconsistency because the loaded id tracker
         // is only used as part of maybe_delete, which is idempotent.
-        let id_tracker = VectorStaticIdTracker(StaticIdTracker::load_from_path(
-            id_tracker_path,
-            deleted_bitset.clone(),
-        )?);
+        let id_tracker = VectorStaticIdTracker {
+            id_tracker: StaticIdTracker::load_from_path(id_tracker_path)?,
+            deleted_bitset: deleted.clone(),
+        };
 
-        Ok(Self::new(original, id_tracker, deleted_bitset))
+        Ok(Self::new(original, id_tracker, deleted))
     }
 
     pub async fn upload_deleted_bitset(

crates/search/src/fragmented_segment.rs

@@ -9,7 +9,10 @@ use tantivy::{
     IndexBuilder,
     SingleSegmentIndexWriter,
 };
-use text_search::tracker::MemoryIdAndDeletionTracker;
+use text_search::tracker::{
+    MemoryDeletionTracker,
+    SearchMemoryIdTracker,
+};
 
 use crate::{
     constants::CONVEX_EN_TOKENIZER,
@@ -24,7 +27,7 @@ const SEGMENT_MAX_SIZE_BYTES: usize = 10_000_000;
 #[allow(dead_code)]
 pub(crate) const ID_TRACKER_PATH: &str = "id_tracker";
 #[allow(dead_code)]
-pub(crate) const DELETED_TANTIVY_IDS_PATH: &str = "deleted_tantivy_ids";
+pub(crate) const ALIVE_BITSET_PATH: &str = "tantivy_alive_bitset";
 #[allow(dead_code)]
 pub(crate) const DELETED_TERMS_PATH: &str = "deleted_terms";
 
@@ -42,7 +45,7 @@ pub async fn build_index(
         .tokenizers()
         .register(CONVEX_EN_TOKENIZER, convex_en());
     let mut segment_writer = SingleSegmentIndexWriter::new(index, SEGMENT_MAX_SIZE_BYTES)?;
-    let mut tracker = MemoryIdAndDeletionTracker::default();
+    let mut id_tracker = SearchMemoryIdTracker::default();
     futures::pin_mut!(revision_stream);
     // Keep track of the document IDs we've seen so we can check for duplicates.
     // We'll discard revisions to documents that we've already seen because we are
@@ -58,13 +61,14 @@ pub async fn build_index(
             let tantivy_document =
                 tantivy_schema.index_into_tantivy_document(new_document, revision_pair.ts());
             let doc_id = segment_writer.add_document(tantivy_document)?;
-            tracker.set_link(convex_id, doc_id)?;
+            id_tracker.set_link(convex_id, doc_id)?;
         }
     }
     segment_writer.finalize()?;
+    id_tracker.write(dir.to_path_buf().join(ID_TRACKER_PATH))?;
+    let tracker = MemoryDeletionTracker::new(document_ids_seen.len() as u32);
     tracker.write(
-        dir.to_path_buf().join(ID_TRACKER_PATH),
-        dir.to_path_buf().join(DELETED_TANTIVY_IDS_PATH),
+        dir.to_path_buf().join(ALIVE_BITSET_PATH),
         dir.to_path_buf().join(DELETED_TERMS_PATH),
     )?;
     Ok(())