[BUG] Handle uninitialized spann segment reader (chroma-core#4348)

Author: warpbuild-benchmark-bot[bot]

rust/worker/chroma_config.yaml

@@ -53,7 +53,7 @@ query_service:
     blockfile_provider:
         arrow:
             block_manager_config:
-                max_block_size_bytes: 16384
+                max_block_size_bytes: 8388608 # 8MB
                 block_cache_config:
                     disk:
                         dir: "/cache/chroma/query-service/block-cache"
@@ -135,7 +135,7 @@ compaction_service:
     blockfile_provider:
         arrow:
             block_manager_config:
-                max_block_size_bytes: 16384
+                max_block_size_bytes: 8388608 # 8MB
                 block_cache_config:
                     lru:
                         capacity: 1000

rust/worker/src/execution/orchestration/spann_knn.rs

@@ -1,6 +1,9 @@
 use async_trait::async_trait;
 use chroma_distance::{normalize, DistanceFunction};
-use chroma_segment::{distributed_spann::SpannSegmentReader, spann_provider::SpannProvider};
+use chroma_segment::{
+    distributed_spann::{SpannSegmentReader, SpannSegmentReaderError},
+    spann_provider::SpannProvider,
+};
 use chroma_system::{
     wrap, ComponentContext, ComponentHandle, Dispatcher, Handler, Orchestrator, TaskMessage,
     TaskResult,
@@ -131,25 +134,6 @@ impl SpannKnnOrchestrator {
             self.send(task, ctx).await;
         }
     }
-
-    async fn set_spann_reader(&mut self, ctx: &ComponentContext<Self>) {
-        let reader_res = SpannSegmentReader::from_segment(
-            &self.collection,
-            &self.knn_filter_output.vector_segment,
-            &self.spann_provider.blockfile_provider,
-            &self.spann_provider.hnsw_provider,
-            self.knn_filter_output.dimension,
-        )
-        .await;
-        let reader = match self.ok_or_terminate(reader_res, ctx).await {
-            Some(reader) => reader,
-            None => {
-                tracing::error!("Failed to create SpannSegmentReader");
-                return;
-            }
-        };
-        self.spann_reader = Some(reader);
-    }
 }
 
 #[async_trait]
@@ -176,16 +160,42 @@ impl Orchestrator for SpannKnnOrchestrator {
             ctx.receiver(),
         );
         tasks.push(knn_log_task);
-        self.set_spann_reader(ctx).await;
-        let head_search_task = wrap(
-            Box::new(self.head_search.clone()),
-            SpannCentersSearchInput {
-                reader: self.spann_reader.clone(),
-                normalized_query: self.normalized_query_emb.clone(),
+        let reader_res = SpannSegmentReader::from_segment(
+            &self.collection,
+            &self.knn_filter_output.vector_segment,
+            &self.spann_provider.blockfile_provider,
+            &self.spann_provider.hnsw_provider,
+            self.knn_filter_output.dimension,
+        )
+        .await;
+        match reader_res {
+            Ok(reader) => {
+                self.spann_reader = Some(reader.clone());
+                // Spawn the centers search task if reader is found.
+                let head_search_task = wrap(
+                    Box::new(self.head_search.clone()),
+                    SpannCentersSearchInput {
+                        reader: Some(reader),
+                        normalized_query: self.normalized_query_emb.clone(),
+                    },
+                    ctx.receiver(),
+                );
+                tasks.push(head_search_task);
+            }
+            Err(e) => match e {
+                // Segment uninited means no compaction yet.
+                SpannSegmentReaderError::UninitializedSegment => {
+                    // If the segment is uninitialized, we can skip the head search.
+                    self.spann_reader = None;
+                    self.heads_searched = true;
+                }
+                _ => {
+                    let _: Option<()> = self
+                        .ok_or_terminate(Err(KnnError::SpannSegmentReaderCreationError(e)), ctx)
+                        .await;
+                }
             },
-            ctx.receiver(),
-        );
-        tasks.push(head_search_task);
+        }
 
         let prefetch_task = wrap(
             Box::new(PrefetchSegmentOperator::new()),

rust/worker/tilt_config.yaml

@@ -52,7 +52,7 @@ query_service:
     blockfile_provider:
         arrow:
             block_manager_config:
-                max_block_size_bytes: 16384
+                max_block_size_bytes: 8388608 # 8MB
                 block_cache_config:
                     disk:
                         dir: "/cache/chroma/query-service/block-cache"
@@ -138,7 +138,7 @@ compaction_service:
     blockfile_provider:
         arrow:
             block_manager_config:
-                max_block_size_bytes: 16384
+                max_block_size_bytes: 8388608 # 8MB
                 block_cache_config:
                     lru:
                         capacity: 1000