Get basic multi segment text searches to work (#26152)

Fixes a couple small issues in the search path:

* We errored if any term was missing from a segment
* Individual files we're being loaded from directories

GitOrigin-RevId: 41593eedb68a8f0afe77de7fdb3d9661bbd84414

Author: sjudd

crates/database/src/tests/text_test_utils.rs

@@ -13,19 +13,33 @@ use common::{
         IndexMetadata,
         TabletIndexMetadata,
     },
-    document::ParsedDocument,
+    document::{
+        ParsedDocument,
+        ResolvedDocument,
+    },
     persistence::PersistenceReader,
+    query::{
+        Query,
+        QueryOperator,
+        QuerySource,
+        Search,
+        SearchFilterExpression,
+    },
     runtime::testing::TestRuntime,
     types::{
         GenericIndexName,
         IndexId,
         IndexName,
         TabletIndexName,
     },
+    version::MIN_NPM_VERSION_FOR_FUZZY_SEARCH,
 };
 use maplit::btreeset;
 use must_let::must_let;
-use search::searcher::InProcessSearcher;
+use search::{
+    searcher::InProcessSearcher,
+    MAX_CANDIDATE_REVISIONS,
+};
 use storage::Storage;
 use sync_types::Timestamp;
 use value::{
@@ -47,6 +61,7 @@ use crate::{
     },
     Database,
     IndexModel,
+    ResolvedQuery,
     TestFacingModel,
     TextIndexFlusher,
     Transaction,
@@ -187,16 +202,62 @@ impl TextFixtures {
         Ok(segments.clone())
     }
 
-    pub async fn add_document(&self, text: &str) -> anyhow::Result<()> {
+    pub async fn add_document(&self, text: &str) -> anyhow::Result<ResolvedDocumentId> {
         let table_name = TABLE_NAME.parse::<TableName>()?;
         let mut tx = self.db.begin_system().await?;
-        add_document(&mut tx, &table_name, text).await?;
+        let doc_id = add_document(&mut tx, &table_name, text).await?;
+        self.db.commit(tx).await?;
+        Ok(doc_id)
+    }
+
+    pub async fn enable_index(
+        &self,
+        index_name: &GenericIndexName<TableName>,
+    ) -> anyhow::Result<()> {
+        let mut tx = self.db.begin_system().await?;
+        IndexModel::new(&mut tx)
+            .enable_index_for_testing(self.namespace, index_name)
+            .await?;
         self.db.commit(tx).await?;
         Ok(())
     }
+
+    pub async fn search(
+        &self,
+        index_name: GenericIndexName<TableName>,
+        query_string: &str,
+    ) -> anyhow::Result<Vec<ResolvedDocument>> {
+        let mut tx = self.db.begin_system().await?;
+        let filters = vec![SearchFilterExpression::Search(
+            SEARCH_FIELD.parse()?,
+            query_string.into(),
+        )];
+        let search = Search {
+            table: index_name.table().clone(),
+            index_name,
+            filters,
+        };
+
+        let query = Query {
+            source: QuerySource::Search(search),
+            operators: vec![QueryOperator::Limit(MAX_CANDIDATE_REVISIONS)],
+        };
+        let mut query_stream = ResolvedQuery::new_with_version(
+            &mut tx,
+            TableNamespace::Global,
+            query,
+            Some(MIN_NPM_VERSION_FOR_FUZZY_SEARCH.clone()),
+        )?;
+        let mut values = vec![];
+        while let Some(value) = query_stream.next(&mut tx, None).await? {
+            values.push(value);
+        }
+        Ok(values)
+    }
 }
 
 const TABLE_NAME: &str = "table";
+const SEARCH_FIELD: &str = "text";
 
 pub struct IndexData {
     pub index_id: IndexId,
@@ -208,7 +269,7 @@ pub struct IndexData {
 pub fn backfilling_text_index() -> anyhow::Result<IndexMetadata<TableName>> {
     let table_name: TableName = TABLE_NAME.parse()?;
     let index_name = IndexName::new(table_name, "search_index".parse()?)?;
-    let search_field: FieldPath = "text".parse()?;
+    let search_field: FieldPath = SEARCH_FIELD.parse()?;
     let filter_field: FieldPath = "channel".parse()?;
     let metadata = IndexMetadata::new_backfilling_search_index(
         index_name,

crates/database/src/text_index_worker/flusher2.rs

@@ -232,7 +232,7 @@ impl<RT: Runtime> TextIndexFlusher2<RT> {
                 TextIndexState::Backfilled(TextIndexSnapshot {
                     data: TextIndexSnapshotData::MultiSegment(segments),
                     ts: backfill_ts,
-                    version: TextSnapshotVersion::V0,
+                    version: TextSnapshotVersion::V2UseStringIds,
                 })
             } else {
                 let cursor = if let Some(cursor) = backfill_result.new_cursor {
@@ -345,10 +345,12 @@ mod tests {
         rt: TestRuntime,
     ) -> anyhow::Result<()> {
         let fixtures = TextFixtures::new(rt).await?;
-        let index = fixtures.insert_backfilling_text_index().await?;
+        let index = fixtures
+            .insert_backfilling_text_index_with_document()
+            .await?;
         let mut flusher = fixtures.new_search_flusher2();
         flusher.step().await?;
-        fixtures.assert_backfilled(&index.name).await?;
+        fixtures.assert_backfilled(&index.index_name).await?;
         Ok(())
     }
 
@@ -367,6 +369,21 @@ mod tests {
         Ok(())
     }
 
+    #[convex_macro::test_runtime]
+    async fn backfill_with_one_document_writes_document(rt: TestRuntime) -> anyhow::Result<()> {
+        let fixtures = TextFixtures::new(rt).await?;
+        let index = fixtures.insert_backfilling_text_index().await?;
+        let doc_id = fixtures.add_document("cat").await?;
+        let mut flusher = fixtures.new_search_flusher2();
+        flusher.step().await?;
+        fixtures.enable_index(&index.name).await?;
+
+        let results = fixtures.search(index.name, "cat").await?;
+        assert_eq!(results.first().unwrap().id(), doc_id);
+
+        Ok(())
+    }
+
     #[convex_macro::test_runtime]
     async fn backfill_with_two_documents_0_max_segment_size_creates_two_segments(
         rt: TestRuntime,
@@ -390,4 +407,34 @@ mod tests {
         assert_eq!(segments.len(), 2);
         Ok(())
     }
+
+    #[convex_macro::test_runtime]
+    async fn backfill_with_two_documents_0_max_segment_size_includes_both_documents(
+        rt: TestRuntime,
+    ) -> anyhow::Result<()> {
+        let fixtures = TextFixtures::new(rt).await?;
+        let IndexMetadata { name, .. } = fixtures.insert_backfilling_text_index().await?;
+
+        let cat_doc_id = fixtures.add_document("cat").await?;
+        let dog_doc_id = fixtures.add_document("dog").await?;
+
+        let mut flusher = fixtures
+            .new_search_flusher_builder()
+            .set_incremental_multipart_threshold_bytes(0)
+            .build();
+        // Build the first segment, which stops because the document size is > 0
+        flusher.step().await?;
+        // Build the second segment and finalize the index metadata.
+        flusher.step().await?;
+
+        fixtures.enable_index(&name).await?;
+
+        let cat_results = fixtures.search(name.clone(), "cat").await?;
+        assert_eq!(cat_results.first().unwrap().id(), cat_doc_id);
+
+        let dog_results = fixtures.search(name, "dog").await?;
+        assert_eq!(dog_results.first().unwrap().id(), dog_doc_id);
+
+        Ok(())
+    }
 }

crates/search/src/archive/cache.rs

@@ -35,6 +35,7 @@ use futures::{
     StreamExt,
     TryStreamExt,
 };
+use itertools::Itertools;
 use storage::{
     Storage,
     StorageCacheKey,
@@ -206,9 +207,9 @@ impl<RT: Runtime> ArchiveFetcher<RT> {
         let bytes_copied = tokio::io::copy_buf(&mut reader, &mut file).await?;
         file.shutdown().await?;
 
-        // We're expecting that the uncompressed tar and its contents are rougnly the
+        // We're expecting that the uncompressed tar and its contents are roughly the
         // same size. There is some file moving / copying going on in
-        // this method, but hopefully it's small enough to be arounding
+        // this method, but hopefully it's small enough to be a rounding
         // error relative to the overall segment size.
         let path = Self::unpack_fragmented_segment_tar(output_file)?;
 
@@ -327,6 +328,34 @@ impl<RT: Runtime> ArchiveCacheManager<RT> {
         result
     }
 
+    pub async fn get_single_file(
+        &self,
+        search_storage: Arc<dyn Storage>,
+        storage_path: &ObjectKey,
+        file_type: SearchFileType,
+    ) -> anyhow::Result<PathBuf> {
+        // The archive cache always dumps things into directories, but we want a
+        // specific file path.
+        let parent_dir: PathBuf = self.get(search_storage, storage_path, file_type).await?;
+        // tokio's async read_dir method punts to a thread pool too, but by using our
+        // own, we can be runtime agnostic.
+        let path = self
+            .blocking_thread_pool
+            .execute(move || try {
+                let paths: Vec<_> = std::fs::read_dir(parent_dir)?
+                    .map_ok(|value| value.path())
+                    .try_collect()?;
+                anyhow::ensure!(
+                    paths.len() == 1,
+                    "Expected one file but found multiple paths: {:?}",
+                    paths,
+                );
+                paths.first().unwrap().to_owned()
+            })
+            .await??;
+        Ok(path)
+    }
+
     async fn get_logged(
         &self,
         search_storage: Arc<dyn Storage>,

crates/search/src/fragmented_segment.rs

@@ -1,7 +1,4 @@
-use std::{
-    path::PathBuf,
-    sync::Arc,
-};
+use std::sync::Arc;
 
 use common::{
     bootstrap_model::index::vector_index::FragmentedVectorSegment,
@@ -65,7 +62,6 @@ use crate::{
 #[derive(Clone)]
 pub(crate) struct FragmentedSegmentFetcher<RT: Runtime> {
     archive_cache: ArchiveCacheManager<RT>,
-    blocking_thread_pool: BoundedThreadPool<RT>,
 }
 
 pub struct FragmentedSegmentStorageKeys {
@@ -77,14 +73,8 @@ pub struct FragmentedSegmentStorageKeys {
 impl<RT: Runtime> FragmentedSegmentFetcher<RT> {
     /// blocking_thread_pool is used for small / fast IO operations and should
     /// be large.
-    pub(crate) fn new(
-        archive_cache: ArchiveCacheManager<RT>,
-        blocking_thread_pool: BoundedThreadPool<RT>,
-    ) -> FragmentedSegmentFetcher<RT> {
-        Self {
-            archive_cache,
-            blocking_thread_pool,
-        }
+    pub(crate) fn new(archive_cache: ArchiveCacheManager<RT>) -> FragmentedSegmentFetcher<RT> {
+        Self { archive_cache }
     }
 
     /// Fetch all parts of all fragmented segments with limited concurrency.
@@ -125,14 +115,14 @@ impl<RT: Runtime> FragmentedSegmentFetcher<RT> {
             SearchFileType::FragmentedVectorSegment,
         );
 
-        let fetch_id_tracker = self.fetch_single_file(
+        let fetch_id_tracker = archive_cache.get_single_file(
             search_storage.clone(),
-            paths.id_tracker,
+            &paths.id_tracker,
             SearchFileType::VectorIdTracker,
         );
-        let fetch_bitset = self.fetch_single_file(
+        let fetch_bitset = archive_cache.get_single_file(
             search_storage.clone(),
-            paths.deleted_bitset,
+            &paths.deleted_bitset,
             SearchFileType::VectorDeletedBitset,
         );
         let (segment, id_tracker, bitset) =
@@ -141,37 +131,6 @@ impl<RT: Runtime> FragmentedSegmentFetcher<RT> {
             segment, id_tracker, bitset,
         ))
     }
-
-    async fn fetch_single_file(
-        &self,
-        search_storage: Arc<dyn Storage>,
-        storage_path: ObjectKey,
-        file_type: SearchFileType,
-    ) -> anyhow::Result<PathBuf> {
-        // The archive cache always dumps things into directories, but we want a
-        // specific file path.
-        let parent_dir: PathBuf = self
-            .archive_cache
-            .get(search_storage, &storage_path, file_type)
-            .await?;
-        // tokio's async read_dir method punts to a thread pool too, but by using our
-        // own, we can be runtime agnostic.
-        let path = self
-            .blocking_thread_pool
-            .execute(move || try {
-                let paths: Vec<_> = std::fs::read_dir(parent_dir)?
-                    .map_ok(|value| value.path())
-                    .try_collect()?;
-                anyhow::ensure!(
-                    paths.len() == 1,
-                    "Expected one file but found multiple paths: {:?}",
-                    paths,
-                );
-                paths.first().unwrap().to_owned()
-            })
-            .await??;
-        Ok(path)
-    }
 }
 
 pub(crate) struct FragmentedSegmentCompactor<RT: Runtime> {