Generalize VectorIndexCompactor (#26367)

We'll use this same basic pattern for text search to start with. Still more to extract which will be done in follow-up PRs

GitOrigin-RevId: 030f97251a6e064f51f4d8516994f9fca90fc84a

Author: sjudd

crates/database/src/index_workers/index_meta.rs

@@ -26,7 +26,10 @@ use common::{
         ObjectKey,
     },
 };
-use search::metrics::SearchType;
+use search::{
+    metrics::SearchType,
+    Searcher,
+};
 use storage::Storage;
 use sync_types::Timestamp;
 use value::{
@@ -50,9 +53,9 @@ pub trait PreviousSegmentsType: Send {
 pub trait SegmentType<T: SearchIndex> {
     fn id(&self) -> &str;
 
-    fn num_deleted(&self) -> u64;
-
     fn statistics(&self) -> anyhow::Result<T::Statistics>;
+
+    fn total_size_bytes(&self, config: &T::DeveloperConfig) -> anyhow::Result<u64>;
 }
 
 #[async_trait]
@@ -128,6 +131,13 @@ pub trait SearchIndex: Clone + Debug {
         storage: Arc<dyn Storage>,
         segments: Self::PreviousSegments,
     ) -> anyhow::Result<Vec<Self::Segment>>;
+
+    async fn execute_compaction(
+        searcher: Arc<dyn Searcher>,
+        search_storage: Arc<dyn Storage>,
+        config: &Self::DeveloperConfig,
+        segments: &Vec<&Self::Segment>,
+    ) -> anyhow::Result<Self::Segment>;
 }
 
 pub trait SegmentStatistics: Default + Debug {
@@ -136,6 +146,10 @@ pub trait SegmentStatistics: Default + Debug {
     fn num_documents(&self) -> u64;
 
     fn num_non_deleted_documents(&self) -> u64;
+
+    fn num_deleted_documents(&self) -> u64 {
+        self.num_documents() - self.num_non_deleted_documents()
+    }
 }
 pub struct SearchIndexConfig<T: SearchIndex> {
     pub developer_config: T::DeveloperConfig,

crates/database/src/index_workers/search_worker.rs

@@ -40,19 +40,21 @@ use crate::{
         TextIndexFlusher2,
     },
     vector_index_worker::{
-        compactor::CompactionConfig,
+        compactor::{
+            CompactionConfig,
+            VectorIndexCompactor2,
+        },
         flusher::new_vector_flusher,
     },
     Database,
     TextIndexFlusher,
-    VectorIndexCompactor,
     VectorIndexFlusher,
 };
 
 /// Builds and compacts text/vector search indexes.
 pub enum SearchIndexWorker<RT: Runtime> {
     VectorFlusher(VectorIndexFlusher<RT>),
-    VectorCompactor(VectorIndexCompactor<RT>),
+    VectorCompactor(VectorIndexCompactor2<RT>),
     TextFlusher(TextIndexFlusher<RT>),
     TextFlusher2(TextIndexFlusher2<RT>),
 }
@@ -103,7 +105,7 @@ impl<RT: Runtime> SearchIndexWorker<RT> {
             runtime.clone(),
             database.clone(),
             Duration::ZERO,
-            SearchIndexWorker::VectorCompactor(VectorIndexCompactor::new(
+            SearchIndexWorker::VectorCompactor(VectorIndexCompactor2::new(
                 database.clone(),
                 searcher,
                 search_storage.clone(),

crates/database/src/index_workers/writer.rs

@@ -44,6 +44,7 @@ use crate::{
             SearchIndex,
             SearchOnDiskState,
             SearchSnapshot,
+            SegmentStatistics,
             SegmentType,
             SnapshotData,
         },
@@ -244,7 +245,9 @@ impl<RT: Runtime, T: SearchIndex> Inner<RT, T> {
             // which creates a new segment with a new id. So if the number of deletes has
             // changed, it's due to an increase from a conflicting write by the
             // flusher.
-            if current_version.num_deleted() != original_segment.num_deleted() {
+            if current_version.statistics()?.num_deleted_documents()
+                != original_segment.statistics()?.num_deleted_documents()
+            {
                 return Ok(true);
             }
         }

crates/database/src/lib.rs

@@ -16,6 +16,7 @@
 #![feature(trait_upcasting)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(cow_is_borrowed)]
+#![feature(try_find)]
 
 mod bootstrap_model;
 mod committer;

crates/database/src/metrics.rs

@@ -1048,7 +1048,7 @@ pub mod vector {
         VECTOR_COMPACTION_COMPACTED_SEGMENT_NUM_VECTORS_TOTAL,
         "Size of the newly generated compacted segment",
     );
-    pub fn log_vector_compaction_compacted_segment_num_vectors_total(total_vectors: u32) {
+    pub fn log_vector_compaction_compacted_segment_num_vectors_total(total_vectors: u64) {
         log_distribution(
             &VECTOR_COMPACTION_COMPACTED_SEGMENT_NUM_VECTORS_TOTAL,
             total_vectors as f64,

crates/database/src/tests/vector_test_utils.rs

@@ -90,7 +90,10 @@ use crate::{
     index_workers::search_flusher::FLUSH_RUNNING_LABEL,
     test_helpers::DbFixturesArgs,
     vector_index_worker::{
-        compactor::CompactionConfig,
+        compactor::{
+            CompactionConfig,
+            VectorIndexCompactor2,
+        },
         flusher::{
             backfill_vector_indexes,
             new_vector_flusher_for_tests,
@@ -102,7 +105,6 @@ use crate::{
     TestFacingModel,
     Transaction,
     UserFacingModel,
-    VectorIndexCompactor,
 };
 
 pub struct VectorFixtures {
@@ -219,15 +221,15 @@ impl VectorFixtures {
         Ok(result)
     }
 
-    pub async fn new_compactor(&self) -> anyhow::Result<VectorIndexCompactor<TestRuntime>> {
+    pub async fn new_compactor(&self) -> anyhow::Result<VectorIndexCompactor2<TestRuntime>> {
         self.new_compactor_with_searchlight(self.searcher.clone())
             .await
     }
 
     pub async fn new_compactor_delete_on_compact(
         &self,
         id_to_delete: ResolvedDocumentId,
-    ) -> anyhow::Result<VectorIndexCompactor<TestRuntime>> {
+    ) -> anyhow::Result<VectorIndexCompactor2<TestRuntime>> {
         let searcher = DeleteOnCompactSearchlight {
             rt: self.rt.clone(),
             db: self.db.clone(),
@@ -244,8 +246,8 @@ impl VectorFixtures {
     async fn new_compactor_with_searchlight(
         &self,
         searcher: Arc<dyn Searcher>,
-    ) -> anyhow::Result<VectorIndexCompactor<TestRuntime>> {
-        Ok(VectorIndexCompactor::new_for_tests(
+    ) -> anyhow::Result<VectorIndexCompactor2<TestRuntime>> {
+        Ok(VectorIndexCompactor2::new_for_tests(
             self.rt.clone(),
             self.db.clone(),
             self.storage.clone(),

crates/database/src/tests/vector_tests.rs

@@ -86,15 +86,17 @@ use crate::{
         IndexData,
         VectorFixtures,
     },
-    vector_index_worker::flusher::{
-        backfill_vector_indexes,
-        new_vector_flusher_for_tests,
+    vector_index_worker::{
+        compactor::VectorIndexCompactor2,
+        flusher::{
+            backfill_vector_indexes,
+            new_vector_flusher_for_tests,
+        },
     },
     Database,
     IndexModel,
     TableModel,
     UserFacingModel,
-    VectorIndexCompactor,
 };
 
 const TABLE_NAME: &str = "test";
@@ -206,7 +208,7 @@ impl<RT: Runtime> Scenario<RT> {
     }
 
     async fn compact(&mut self) -> anyhow::Result<()> {
-        VectorIndexCompactor::process_all_in_test(
+        VectorIndexCompactor2::process_all_in_test(
             self.rt.clone(),
             self.database.clone(),
             self.search_storage.clone(),

crates/database/src/text_index_worker/text_meta.rs

@@ -43,6 +43,7 @@ use search::{
     searcher::SegmentTermMetadataFetcher,
     NewTextSegment,
     PreviousTextSegments,
+    Searcher,
     TantivySearchIndexSchema,
     UpdatableTextSegment,
 };
@@ -110,16 +111,19 @@ impl SegmentType<TextSearchIndex> for FragmentedTextSegment {
         &self.id
     }
 
-    fn num_deleted(&self) -> u64 {
-        self.num_deleted_documents
-    }
-
     fn statistics(&self) -> anyhow::Result<TextStatistics> {
         Ok(TextStatistics {
             num_indexed_documents: self.num_indexed_documents,
             num_deleted_documents: self.num_deleted_documents,
         })
     }
+
+    fn total_size_bytes(
+        &self,
+        _config: &<TextSearchIndex as SearchIndex>::DeveloperConfig,
+    ) -> anyhow::Result<u64> {
+        Ok(self.size_bytes_total)
+    }
 }
 #[derive(Clone)]
 pub struct BuildTextIndexArgs {
@@ -262,6 +266,15 @@ impl SearchIndex for TextSearchIndex {
     fn search_type() -> SearchType {
         SearchType::Text
     }
+
+    async fn execute_compaction(
+        _searcher: Arc<dyn Searcher>,
+        _search_storage: Arc<dyn Storage>,
+        _config: &Self::DeveloperConfig,
+        _segments: &Vec<&Self::Segment>,
+    ) -> anyhow::Result<Self::Segment> {
+        unimplemented!()
+    }
 }
 
 #[derive(Debug, Default)]