Use search2 query path with a single segment (#26037)

emmaling27 · Convex, Inc. · commit aaa82ef6098d · 2024-05-21T22:20:11.000Z
This PR adds a knob for using the `search2` query path and adds support for using `search2` with the existing single-segment search indexes. For the single search index, the id tracker is optional and we construct an empty deletion tracker.

GitOrigin-RevId: 03734183805da31a713217fd7250e32cfa75e64c
diff --git a/crates/common/src/knobs.rs b/crates/common/src/knobs.rs
@@ -1018,3 +1018,7 @@ pub static BUILD_MULTI_SEGMENT_TEXT_INDEXES: LazyLock<bool> =
 /// the "backend_startup" domain keyed by db cluster name.
 pub static STARTUP_RATE_LIMIT_ENABLED: LazyLock<bool> =
     LazyLock::new(|| env_config("STARTUP_RATE_LIMIT_ENABLED", false));
+
+/// Use multi segment search algorithm for search queries.
+pub static USE_MULTI_SEGMENT_SEARCH_QUERY: LazyLock<bool> =
+    LazyLock::new(|| env_config("USE_MULTI_SEGMENT_SEARCH_QUERY", false));
diff --git a/crates/database/src/tests/randomized_search_tests.rs b/crates/database/src/tests/randomized_search_tests.rs
@@ -71,11 +71,11 @@ use search::{
     scoring::Bm25StatisticsDiff,
     searcher::{
         Bm25Stats,
-        FragmentedTextSegmentStorageKeys,
         InProcessSearcher,
         PostingListMatch,
         PostingListQuery,
         Term,
+        TextStorageKeys,
         TokenMatch,
         TokenQuery,
     },
@@ -963,10 +963,18 @@ impl Searcher for BrokenSearcher {
         anyhow::bail!("要");
     }
 
+    async fn number_of_segments(
+        &self,
+        _search_storage: Arc<dyn Storage>,
+        _storage_key: ObjectKey,
+    ) -> anyhow::Result<usize> {
+        anyhow::bail!("wut")
+    }
+
     async fn query_tokens(
         &self,
         _: Arc<dyn Storage>,
-        _: FragmentedTextSegmentStorageKeys,
+        _: TextStorageKeys,
         _: Vec<TokenQuery>,
         _: usize,
     ) -> anyhow::Result<Vec<TokenMatch>> {
@@ -976,7 +984,7 @@ impl Searcher for BrokenSearcher {
     async fn query_bm25_stats(
         &self,
         _: Arc<dyn Storage>,
-        _: FragmentedTextSegmentStorageKeys,
+        _: TextStorageKeys,
         _: Vec<Term>,
     ) -> anyhow::Result<Bm25Stats> {
         anyhow::bail!("plein")
@@ -985,7 +993,7 @@ impl Searcher for BrokenSearcher {
     async fn query_posting_lists(
         &self,
         _: Arc<dyn Storage>,
-        _: FragmentedTextSegmentStorageKeys,
+        _: TextStorageKeys,
         _: PostingListQuery,
     ) -> anyhow::Result<Vec<PostingListMatch>> {
         anyhow::bail!("texte");
diff --git a/crates/database/src/tests/vector_test_utils.rs b/crates/database/src/tests/vector_test_utils.rs
@@ -51,12 +51,12 @@ use search::{
     scoring::Bm25StatisticsDiff,
     searcher::{
         Bm25Stats,
-        FragmentedTextSegmentStorageKeys,
         InProcessSearcher,
         PostingListMatch,
         PostingListQuery,
         Searcher,
         Term,
+        TextStorageKeys,
         TokenMatch,
         TokenQuery,
     },
@@ -576,10 +576,20 @@ impl<RT: Runtime> Searcher for DeleteOnCompactSearchlight<RT> {
             .await
     }
 
+    async fn number_of_segments(
+        &self,
+        search_storage: Arc<dyn Storage>,
+        storage_key: ObjectKey,
+    ) -> anyhow::Result<usize> {
+        self.searcher
+            .number_of_segments(search_storage, storage_key)
+            .await
+    }
+
     async fn query_tokens(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         queries: Vec<TokenQuery>,
         max_results: usize,
     ) -> anyhow::Result<Vec<TokenMatch>> {
@@ -591,7 +601,7 @@ impl<RT: Runtime> Searcher for DeleteOnCompactSearchlight<RT> {
     async fn query_bm25_stats(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         terms: Vec<Term>,
     ) -> anyhow::Result<Bm25Stats> {
         self.searcher
@@ -602,7 +612,7 @@ impl<RT: Runtime> Searcher for DeleteOnCompactSearchlight<RT> {
     async fn query_posting_lists(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         query: PostingListQuery,
     ) -> anyhow::Result<Vec<PostingListMatch>> {
         self.searcher
diff --git a/crates/pb/protos/searchlight.proto b/crates/pb/protos/searchlight.proto
@@ -11,6 +11,8 @@ service Searchlight {
   rpc ExecuteVectorCompaction(VectorCompactionRequest) returns (VectorCompactionResponse);
   rpc QueueVectorPrefetch(VectorPrefetchRequest) returns (VectorPrefetchResponse);
 
+  rpc NumberOfSegments(SegmentRequest) returns (SegmentResponse);
+
   // Query a set of tokens against the term dictionary, optionally allowing
   // for fuzzy matching and prefix matching. Take the top `K` results with
   // respect to to `(edit distance, term)` lexicographical order.
@@ -214,6 +216,15 @@ message StorageType {
   }
 }
 
+message SegmentRequest {
+    StorageKey segment = 1;
+    StorageType storage_type = 2;
+}
+
+message SegmentResponse {
+    uint32 number_of_segments = 1;
+}
+
 message QueryTokensRequest {
   StorageType storage_type = 1;
   FragmentedTextSegmentPaths segment = 2;
@@ -223,8 +234,20 @@ message QueryTokensRequest {
 
 message FragmentedTextSegmentPaths {
   StorageKey segment = 1;
-  StorageKey id_tracker = 2;
-  StorageKey deletions = 3;
+  oneof segment_metadata {
+    SingleSegmentMetadata single_segment = 2;
+    MultiSegmentMetadata multi_segment = 3;
+  }
+}
+
+message SingleSegmentMetadata {
+   optional uint32 segment_ord = 1;
+}
+
+message MultiSegmentMetadata {
+  StorageKey id_tracker = 1;
+  StorageKey deleted_terms_table = 2;
+  StorageKey alive_bitset = 3;
 }
 
 message TokenQuery {
diff --git a/crates/search/src/lib.rs b/crates/search/src/lib.rs
@@ -47,6 +47,7 @@ use common::{
     },
     document::ResolvedDocument,
     index::IndexKeyBytes,
+    knobs::USE_MULTI_SEGMENT_SEARCH_QUERY,
     query::{
         search_value_to_bytes,
         InternalSearch,
@@ -84,7 +85,7 @@ use query::{
     RevisionWithKeys,
     TextQueryTerm,
 };
-use searcher::FragmentedTextSegmentStorageKeys;
+use searcher::TextStorageKeys;
 use storage::Storage;
 pub use tantivy::Document as TantivyDocument;
 use tantivy::{
@@ -410,7 +411,7 @@ impl TantivySearchIndexSchema {
         compiled_query: CompiledQuery,
         memory_index: &MemorySearchIndex,
         search_storage: Arc<dyn Storage>,
-        segments: Vec<FragmentedTextSegmentStorageKeys>,
+        segments: Vec<TextStorageKeys>,
         disk_index_ts: Timestamp,
         searcher: Arc<dyn Searcher>,
     ) -> anyhow::Result<RevisionWithKeys> {
@@ -438,11 +439,11 @@ impl TantivySearchIndexSchema {
         // and merge the results to get the top terms.
         let mut match_aggregator = TokenMatchAggregator::new(MAX_UNIQUE_QUERY_TERMS);
         memory_index.query_tokens(&token_queries, &mut match_aggregator)?;
-        for segment in &segments {
+        for segment in segments.clone() {
             let segment_token_matches = searcher
                 .query_tokens(
                     search_storage.clone(),
-                    segment.clone(),
+                    segment,
                     token_queries.clone(),
                     MAX_UNIQUE_QUERY_TERMS,
                 )
@@ -482,9 +483,9 @@ impl TantivySearchIndexSchema {
         // Step 3: Given the terms we decided on, query BM25 statistics across all of
         // the indexes and merge their results.
         let mut bm25_stats = Bm25Stats::empty();
-        for segment in &segments {
+        for segment in segments.clone() {
             bm25_stats += searcher
-                .query_bm25_stats(search_storage.clone(), segment.clone(), terms.clone())
+                .query_bm25_stats(search_storage.clone(), segment, terms.clone())
                 .await?;
         }
         bm25_stats = memory_index.update_bm25_stats(disk_index_ts, &terms, bm25_stats)?;
@@ -545,9 +546,9 @@ impl TantivySearchIndexSchema {
                 &mut match_aggregator,
             )?;
         }
-        for segment in &segments {
+        for segment in segments {
             let segment_matches = searcher
-                .query_posting_lists(search_storage.clone(), segment.clone(), query.clone())
+                .query_posting_lists(search_storage.clone(), segment, query.clone())
                 .await?;
             for m in segment_matches {
                 if !match_aggregator.insert(m) {
@@ -590,6 +591,27 @@ impl TantivySearchIndexSchema {
         disk_index_ts: Timestamp,
         searcher: Arc<dyn Searcher>,
     ) -> anyhow::Result<RevisionWithKeys> {
+        if *USE_MULTI_SEGMENT_SEARCH_QUERY {
+            let number_of_segments = searcher
+                .number_of_segments(search_storage.clone(), disk_index.clone())
+                .await?;
+            let segments = (0..number_of_segments)
+                .map(|i| TextStorageKeys::SingleSegment {
+                    storage_key: disk_index.clone(),
+                    segment_ord: i as u32,
+                })
+                .collect();
+            return self
+                .search2(
+                    compiled_query,
+                    memory_index,
+                    search_storage,
+                    segments,
+                    disk_index_ts,
+                    searcher,
+                )
+                .await;
+        }
         // 1. Fetch the memory index matches for each QueryTerm in the query and bound.
         let (term_shortlist, term_shortlist_ids) =
             memory_index.bound_and_evaluate_query_terms(&compiled_query.text_query);
diff --git a/crates/search/src/search_index_manager.rs b/crates/search/src/search_index_manager.rs
@@ -49,7 +49,10 @@ use crate::{
         CompiledQuery,
         RevisionWithKeys,
     },
-    searcher::FragmentedTextSegmentStorageKeys,
+    searcher::{
+        FragmentedTextStorageKeys,
+        TextStorageKeys,
+    },
     QueryResults,
     Searcher,
     TantivySearchIndexSchema,
@@ -292,7 +295,11 @@ impl SearchIndexManager {
                         segments
                             .iter()
                             .cloned()
-                            .map(FragmentedTextSegmentStorageKeys::from)
+                            .map(|segment| {
+                                TextStorageKeys::MultiSegment(FragmentedTextStorageKeys::from(
+                                    segment,
+                                ))
+                            })
                             .collect(),
                         *disk_index_ts,
                         searcher,
diff --git a/crates/search/src/searcher/in_process.rs b/crates/search/src/searcher/in_process.rs
@@ -25,9 +25,9 @@ use vector::{
 
 use super::searcher::{
     Bm25Stats,
-    FragmentedTextSegmentStorageKeys,
     PostingListMatch,
     PostingListQuery,
+    TextStorageKeys,
     TokenMatch,
     TokenQuery,
 };
@@ -61,10 +61,18 @@ impl Searcher for SearcherStub {
         Ok(SearchQueryResult::empty())
     }
 
+    async fn number_of_segments(
+        &self,
+        _search_storage: Arc<dyn Storage>,
+        _storage_key: ObjectKey,
+    ) -> anyhow::Result<usize> {
+        Ok(1)
+    }
+
     async fn query_tokens(
         &self,
         _search_storage: Arc<dyn Storage>,
-        _storage_keys: FragmentedTextSegmentStorageKeys,
+        _storage_keys: TextStorageKeys,
         _queries: Vec<TokenQuery>,
         _max_results: usize,
     ) -> anyhow::Result<Vec<TokenMatch>> {
@@ -74,7 +82,7 @@ impl Searcher for SearcherStub {
     async fn query_bm25_stats(
         &self,
         _search_storage: Arc<dyn Storage>,
-        _storage_keys: FragmentedTextSegmentStorageKeys,
+        _storage_keys: TextStorageKeys,
         _terms: Vec<Term>,
     ) -> anyhow::Result<Bm25Stats> {
         Ok(Bm25Stats {
@@ -87,7 +95,7 @@ impl Searcher for SearcherStub {
     async fn query_posting_lists(
         &self,
         _search_storage: Arc<dyn Storage>,
-        _storage_keys: FragmentedTextSegmentStorageKeys,
+        _storage_keys: TextStorageKeys,
         _query: PostingListQuery,
     ) -> anyhow::Result<Vec<PostingListMatch>> {
         Ok(vec![])
@@ -161,10 +169,20 @@ impl<RT: Runtime> Searcher for InProcessSearcher<RT> {
             .await
     }
 
+    async fn number_of_segments(
+        &self,
+        search_storage: Arc<dyn Storage>,
+        storage_key: ObjectKey,
+    ) -> anyhow::Result<usize> {
+        self.searcher
+            .number_of_segments(search_storage, storage_key)
+            .await
+    }
+
     async fn query_tokens(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         queries: Vec<TokenQuery>,
         max_results: usize,
     ) -> anyhow::Result<Vec<TokenMatch>> {
@@ -176,7 +194,7 @@ impl<RT: Runtime> Searcher for InProcessSearcher<RT> {
     async fn query_bm25_stats(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         terms: Vec<Term>,
     ) -> anyhow::Result<Bm25Stats> {
         self.searcher
@@ -187,7 +205,7 @@ impl<RT: Runtime> Searcher for InProcessSearcher<RT> {
     async fn query_posting_lists(
         &self,
         search_storage: Arc<dyn Storage>,
-        storage_keys: FragmentedTextSegmentStorageKeys,
+        storage_keys: TextStorageKeys,
         query: PostingListQuery,
     ) -> anyhow::Result<Vec<PostingListMatch>> {
         self.searcher
diff --git a/crates/search/src/searcher/mod.rs b/crates/search/src/searcher/mod.rs
@@ -11,12 +11,13 @@ pub use in_process::{
 };
 pub use searcher::{
     Bm25Stats,
-    FragmentedTextSegmentStorageKeys,
+    FragmentedTextStorageKeys,
     PostingListMatch,
     PostingListQuery,
     Searcher,
     SearcherImpl,
     Term,
+    TextStorageKeys,
     TokenMatch,
     TokenQuery,
 };
diff --git a/crates/search/src/searcher/searcher.rs b/crates/search/src/searcher/searcher.rs
diff --git a/crates/text_search/src/tracker.rs b/crates/text_search/src/tracker.rs
