Implement upload_previous_segments for text indexes (#26033)

GitOrigin-RevId: 1311c504f4ca466145e5f02b5199c81766ca9321

Author: sjudd

crates/database/src/text_index_worker/text_meta.rs

@@ -134,10 +134,16 @@ impl SearchIndex for TextSearchIndex {
     }
 
     async fn upload_previous_segments(
-        _storage: Arc<dyn Storage>,
-        _segments: Self::PreviousSegments,
+        storage: Arc<dyn Storage>,
+        segments: Self::PreviousSegments,
     ) -> anyhow::Result<Vec<Self::Segment>> {
-        anyhow::bail!("Not implemented");
+        segments
+            .0
+            .into_iter()
+            .map(|segment| segment.upload_metadata(storage.clone()))
+            .collect::<FuturesUnordered<_>>()
+            .try_collect::<Vec<_>>()
+            .await
     }
 
     fn estimate_document_size(_schema: &Self::Schema, _doc: &ResolvedDocument) -> u64 {

crates/search/src/fragmented_segment.rs

@@ -320,10 +320,9 @@ impl MutableFragmentedSegmentMetadata {
         let mut buf = vec![];
         self.mutated_deleted_bitset.write(&mut buf)?;
 
-        let filename = format!("deleted_bitset");
         let object_key = upload_single_file(
             &mut buf.as_slice(),
-            filename,
+            "deleted_bitset".to_string(),
             storage.clone(),
             SearchFileType::VectorDeletedBitset,
         )

crates/search/src/incremental_index.rs

@@ -43,7 +43,11 @@ use crate::{
     archive::extract_zip,
     constants::CONVEX_EN_TOKENIZER,
     convex_en,
-    disk_index::download_single_file_zip,
+    disk_index::{
+        download_single_file_zip,
+        upload_single_file,
+    },
+    SearchFileType,
     TantivySearchIndexSchema,
     SEARCH_FIELD_ID,
 };
@@ -68,6 +72,7 @@ pub struct UpdatableTextSegment {
     inverted_index: Arc<InvertedIndexReader>,
     id_tracker: StaticIdTracker,
     deletion_tracker: MemoryDeletionTracker,
+    original: FragmentedTextSegment,
 }
 
 fn inverted_index_from_index(index: &Index) -> anyhow::Result<Arc<InvertedIndexReader>> {
@@ -94,6 +99,51 @@ impl UpdatableTextSegment {
             inverted_index,
             id_tracker,
             deletion_tracker,
+            // TODO(sam): We should probably create this outside of this method, then pass it
+            // through here. For now this is unused in these tests.
+            original: FragmentedTextSegment {
+                segment_key: "segment".try_into()?,
+                id_tracker_key: "id_tracker".try_into()?,
+                deleted_terms_table_key: "deleted_terms".try_into()?,
+                alive_bitset_key: "bitset".try_into()?,
+                num_indexed_documents: 0,
+                id: "test_id".to_string(),
+            },
+        })
+    }
+
+    pub async fn upload_metadata(
+        self,
+        storage: Arc<dyn Storage>,
+    ) -> anyhow::Result<FragmentedTextSegment> {
+        // TODO(CX-6511): Skip the upload and return the original file if this segment
+        // wasn't modified.
+
+        let mut bitset_buf = vec![];
+        let mut deleted_terms_buf = vec![];
+        self.deletion_tracker
+            .write(&mut bitset_buf, &mut deleted_terms_buf)?;
+
+        let mut bitset_slice = bitset_buf.as_slice();
+        let upload_bitset = upload_single_file(
+            &mut bitset_slice,
+            "alive_bitset".to_string(),
+            storage.clone(),
+            SearchFileType::TextAliveBitset,
+        );
+        let mut deleted_terms_slice = deleted_terms_buf.as_slice();
+        let upload_deleted_terms = upload_single_file(
+            &mut deleted_terms_slice,
+            "deleted_terms".to_string(),
+            storage.clone(),
+            SearchFileType::TextDeletedTerms,
+        );
+        let (alive_bitset_key, deleted_terms_table_key) =
+            futures::try_join!(upload_bitset, upload_deleted_terms)?;
+        Ok(FragmentedTextSegment {
+            deleted_terms_table_key,
+            alive_bitset_key,
+            ..self.original
         })
     }
 
@@ -150,6 +200,7 @@ impl UpdatableTextSegment {
             inverted_index,
             id_tracker,
             deletion_tracker,
+            original,
         })
     }
 }
@@ -292,7 +343,7 @@ pub async fn build_new_segment(
     let new_deletion_tracker = MemoryDeletionTracker::new(new_id_tracker.num_ids() as u32);
     let alive_bit_set_path = dir.join(ALIVE_BITSET_PATH);
     let deleted_terms_path = dir.join(DELETED_TERMS_PATH);
-    new_deletion_tracker.write(&alive_bit_set_path, &deleted_terms_path)?;
+    new_deletion_tracker.write_to_path(&alive_bit_set_path, &deleted_terms_path)?;
     let id_tracker_path = dir.join(ID_TRACKER_PATH);
     new_id_tracker.write(&id_tracker_path)?;
 
@@ -371,7 +422,7 @@ pub async fn merge_segments(
     let tracker = MemoryDeletionTracker::new(num_docs as u32);
     let alive_bit_set_path = dir.to_path_buf().join(ALIVE_BITSET_PATH);
     let deleted_terms_path = dir.to_path_buf().join(DELETED_TERMS_PATH);
-    tracker.write(&alive_bit_set_path, &deleted_terms_path)?;
+    tracker.write_to_path(&alive_bit_set_path, &deleted_terms_path)?;
     Ok(TextSegmentPaths {
         index_path: index_dir,
         id_tracker_path,

crates/search/src/searcher/searcher.rs

@@ -1554,7 +1554,7 @@ mod tests {
             let dir = previous_segment_dir.join(format!("segment_{i}"));
             std::fs::create_dir(&dir)?;
             updated_deletion_tracker
-                .write(dir.join(ALIVE_BITSET_PATH), dir.join(DELETED_TERMS_PATH))?;
+                .write_to_path(dir.join(ALIVE_BITSET_PATH), dir.join(DELETED_TERMS_PATH))?;
             previous_segment_dirs.push(dir);
         }
         Ok(TestIndex {

crates/text_search/src/tracker.rs

@@ -301,25 +301,30 @@ impl MemoryDeletionTracker {
         self.num_deleted_terms = num_deleted_terms;
     }
 
-    pub fn write<P: AsRef<Path>>(
+    pub fn write_to_path<P: AsRef<Path>>(
         self,
         alive_bitset_path: P,
         deleted_terms_path: P,
     ) -> anyhow::Result<()> {
-        {
-            let mut out = BufWriter::new(File::create(alive_bitset_path)?);
-            self.alive_bitset.serialize(&mut out)?;
-            out.into_inner()?.sync_all()?;
-        }
-        {
-            let mut out = BufWriter::new(File::create(deleted_terms_path)?);
-            Self::write_deleted_terms(
-                self.term_to_deleted_documents,
-                self.num_deleted_terms,
-                &mut out,
-            )?;
-            out.into_inner()?.sync_all()?;
-        }
+        let mut alive_bitset = BufWriter::new(File::create(alive_bitset_path)?);
+        let mut deleted_terms = BufWriter::new(File::create(deleted_terms_path)?);
+        self.write(&mut alive_bitset, &mut deleted_terms)?;
+        alive_bitset.into_inner()?.sync_all()?;
+        deleted_terms.into_inner()?.sync_all()?;
+        Ok(())
+    }
+
+    pub fn write(
+        self,
+        mut alive_bitset: impl Write,
+        mut deleted_terms: impl Write,
+    ) -> anyhow::Result<()> {
+        self.alive_bitset.serialize(&mut alive_bitset)?;
+        Self::write_deleted_terms(
+            self.term_to_deleted_documents,
+            self.num_deleted_terms,
+            &mut deleted_terms,
+        )?;
         Ok(())
     }