Implement upload_new_segment for text indexes (#26017)

GitOrigin-RevId: b16a4838854541bbc6ab0389cfad685d5fbb8217

Author: sjudd

crates/database/src/text_index_worker/text_meta.rs

@@ -32,6 +32,7 @@ use futures::{
 };
 use search::{
     build_new_segment,
+    disk_index::upload_text_segment,
     PreviousTextSegments,
     TantivySearchIndexSchema,
     TextSegmentPaths,
@@ -166,11 +167,11 @@ impl SearchIndex for TextSearchIndex {
     }
 
     async fn upload_new_segment<RT: Runtime>(
-        _rt: &RT,
-        _storage: Arc<dyn Storage>,
-        _new_segment: Self::NewSegment,
+        rt: &RT,
+        storage: Arc<dyn Storage>,
+        new_segment: Self::NewSegment,
     ) -> anyhow::Result<Self::Segment> {
-        anyhow::bail!("Not implemented")
+        upload_text_segment(rt, storage, new_segment).await
     }
 
     fn segment_id(segment: &Self::Segment) -> String {

crates/database/src/vector_index_worker/vector_meta.rs

@@ -30,7 +30,7 @@ use futures::{
     TryStreamExt,
 };
 use search::{
-    disk_index::upload_segment,
+    disk_index::upload_vector_segment,
     fragmented_segment::MutableFragmentedSegmentMetadata,
 };
 use storage::Storage;
@@ -162,7 +162,7 @@ impl SearchIndex for VectorSearchIndex {
         storage: Arc<dyn Storage>,
         new_segment: Self::NewSegment,
     ) -> anyhow::Result<Self::Segment> {
-        upload_segment(rt, storage, new_segment).await
+        upload_vector_segment(rt, storage, new_segment).await
     }
 
     fn segment_id(segment: &Self::Segment) -> String {

crates/search/src/archive/cache.rs

@@ -378,7 +378,7 @@ fn is_immutable(search_file_type: SearchFileType) -> bool {
         // Text indexes do not appear to be read in readonly mode.
         SearchFileType::Text => false,
         SearchFileType::TextIdTracker => true,
-        SearchFileType::TextDeletedBitset => true,
+        SearchFileType::TextAliveBitset => true,
         SearchFileType::TextDeletedTerms => true,
     }
 }

crates/search/src/disk_index.rs

@@ -23,7 +23,10 @@ use common::{
         FuturesAsyncWriteCompatExt,
         TokioAsyncWriteCompatExt,
     },
-    bootstrap_model::index::vector_index::FragmentedVectorSegment,
+    bootstrap_model::index::{
+        text_index::FragmentedTextSegment,
+        vector_index::FragmentedVectorSegment,
+    },
     runtime::Runtime,
     types::ObjectKey,
 };
@@ -66,6 +69,7 @@ use crate::{
     },
     SearchFileType,
     TantivySearchIndexSchema,
+    TextSegmentPaths,
 };
 
 static SEARCH_INDEXING_MEMORY_ARENA_BYTES: LazyLock<usize> =
@@ -162,7 +166,54 @@ pub async fn download_single_file_zip<P: AsRef<Path>>(
     Ok(())
 }
 
-pub async fn upload_segment<RT: Runtime>(
+pub async fn upload_text_segment<RT: Runtime>(
+    rt: &RT,
+    storage: Arc<dyn Storage>,
+    new_segment: TextSegmentPaths,
+) -> anyhow::Result<FragmentedTextSegment> {
+    let TextSegmentPaths {
+        index_path,
+        id_tracker_path,
+        alive_bit_set_path,
+        deleted_terms_path,
+    } = new_segment;
+    let upload_index =
+        upload_index_archive_from_path(index_path, storage.clone(), SearchFileType::Text);
+    let upload_id_tracker = upload_single_file_from_path(
+        id_tracker_path,
+        storage.clone(),
+        SearchFileType::TextIdTracker,
+    );
+    let upload_bitset = upload_single_file_from_path(
+        alive_bit_set_path,
+        storage.clone(),
+        SearchFileType::TextAliveBitset,
+    );
+    let upload_deleted_terms = upload_single_file_from_path(
+        deleted_terms_path,
+        storage.clone(),
+        SearchFileType::TextDeletedTerms,
+    );
+    let result = futures::try_join!(
+        upload_index,
+        upload_id_tracker,
+        upload_bitset,
+        upload_deleted_terms
+    )?;
+    let (segment_key, id_tracker_key, alive_bitset_key, deleted_terms_table_key) = result;
+    Ok(FragmentedTextSegment {
+        segment_key,
+        id_tracker_key,
+        deleted_terms_table_key,
+        alive_bitset_key,
+        // TODO(sam): Wrap TextSegmentPaths in another struct with some metadata and pass it through
+        // here.
+        num_indexed_documents: 0,
+        id: rt.new_uuid_v4().to_string(),
+    })
+}
+
+pub async fn upload_vector_segment<RT: Runtime>(
     rt: &RT,
     storage: Arc<dyn Storage>,
     new_segment: VectorDiskSegmentValues,

crates/search/src/fragmented_segment.rs

@@ -46,8 +46,8 @@ use crate::{
     archive::cache::ArchiveCacheManager,
     disk_index::{
         download_single_file_zip,
-        upload_segment,
         upload_single_file,
+        upload_vector_segment,
     },
     metrics::{
         log_compacted_segment_size_bytes,
@@ -242,7 +242,7 @@ impl<RT: Runtime> FragmentedSegmentCompactor<RT> {
             })
             .await??;
 
-        let result = upload_segment(&self.rt, search_storage, new_segment).await?;
+        let result = upload_vector_segment(&self.rt, search_storage, new_segment).await?;
         // Ensure we own the temp dir through the entire upload
         drop(tmp_dir);
         tracing::debug!("Compacted {} segments", total_segments);

crates/search/src/lib.rs

@@ -892,7 +892,7 @@ pub enum SearchFileType {
     VectorIdTracker,
     Text,
     TextIdTracker,
-    TextDeletedBitset,
+    TextAliveBitset,
     TextDeletedTerms,
 }
 

crates/search/src/metrics.rs

@@ -426,7 +426,7 @@ impl SearchFileType {
             SearchFileType::VectorIdTracker => "vector_id_tracker",
             SearchFileType::Text => "text",
             SearchFileType::TextIdTracker => "text_id_tracker",
-            SearchFileType::TextDeletedBitset => "text_deleted_bitset",
+            SearchFileType::TextAliveBitset => "text_alive_bitset",
             SearchFileType::TextDeletedTerms => "text_deleted_terms",
             SearchFileType::FragmentedVectorSegment => "fragmented_vector_segment",
         };

crates/search/src/searcher/searcher.rs

@@ -373,7 +373,7 @@ impl<RT: Runtime> Searcher for SearcherImpl<RT> {
             self.archive_cache.get(
                 search_storage.clone(),
                 &storage_keys.segment,
-                SearchFileType::TextDeletedBitset
+                SearchFileType::TextAliveBitset
             ),
             self.archive_cache.get(
                 search_storage.clone(),
@@ -409,7 +409,7 @@ impl<RT: Runtime> Searcher for SearcherImpl<RT> {
             self.archive_cache.get(
                 search_storage.clone(),
                 &storage_keys.segment,
-                SearchFileType::TextDeletedBitset
+                SearchFileType::TextAliveBitset
             ),
             self.archive_cache.get(
                 search_storage.clone(),
@@ -450,7 +450,7 @@ impl<RT: Runtime> Searcher for SearcherImpl<RT> {
             self.archive_cache.get(
                 search_storage.clone(),
                 &storage_keys.segment,
-                SearchFileType::TextDeletedBitset
+                SearchFileType::TextAliveBitset
             ),
             self.archive_cache.get(
                 search_storage.clone(),