[ENH]: Prefetch posting list in query + compactor (#4257)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - ...
 - New functionality
   - Prefetches Posting lists in the beginning of query and compaction. Also, tracks which sparse index ids were already prefetched so as to not do it again. This state is not persistent across restarts and has a ttl.

## Test plan
*How are these changes tested?*
- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
None

Author: sanketkedia

rust/blockstore/src/arrow/provider.rs

@@ -21,7 +21,11 @@ use chroma_storage::{
     admissioncontrolleds3::StorageRequestPriority, GetOptions, PutOptions, Storage,
 };
 use futures::{stream::FuturesUnordered, StreamExt};
-use std::sync::Arc;
+use std::{
+    collections::HashMap,
+    sync::Arc,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
 use thiserror::Error;
 use tracing::{Instrument, Span};
 use uuid::Uuid;
@@ -43,6 +47,8 @@ impl ChromaError for ArrowBlockfileProviderPrefetchError {
     }
 }
 
+const PREFETCH_TTL_HOURS: u64 = 8;
+
 /// A BlockFileProvider that creates ArrowBlockfiles (Arrow-backed blockfiles used for production).
 /// For now, it keeps a simple local cache of blockfiles.
 #[derive(Clone)]
@@ -83,6 +89,9 @@ impl ArrowBlockfileProvider {
     }
 
     pub async fn prefetch(&self, id: &Uuid) -> Result<usize, ArrowBlockfileProviderPrefetchError> {
+        if !self.root_manager.should_prefetch(id) {
+            return Ok(0);
+        }
         // We call .get_all_block_ids() here instead of just reading the root because reading the root requires a concrete Key type.
         let block_ids = self
             .root_manager
@@ -99,10 +108,13 @@ impl ArrowBlockfileProvider {
         }
         let count = futures.len();
 
+        tracing::info!("Prefetching {} blocks for blockfile ID: {:?}", count, id);
+
         while let Some(result) = futures.next().await {
             result?;
         }
 
+        tracing::info!("Prefetched {} blocks for blockfile ID: {:?}", count, id);
         Ok(count)
     }
 
@@ -175,6 +187,7 @@ impl ArrowBlockfileProvider {
     pub async fn clear(&self) -> Result<(), CacheError> {
         self.block_manager.block_cache.clear().await?;
         self.root_manager.cache.clear().await?;
+        self.root_manager.prefetched_roots.lock().clear();
         Ok(())
     }
 }
@@ -466,12 +479,18 @@ impl ChromaError for RootManagerError {
 pub struct RootManager {
     cache: Arc<dyn PersistentCache<Uuid, RootReader>>,
     storage: Storage,
+    // Sparse indexes that have already been prefetched and don't need to be prefetched again.
+    prefetched_roots: Arc<parking_lot::Mutex<HashMap<Uuid, Duration>>>,
 }
 
 impl RootManager {
     pub fn new(storage: Storage, cache: Box<dyn PersistentCache<Uuid, RootReader>>) -> Self {
         let cache: Arc<dyn PersistentCache<Uuid, RootReader>> = cache.into();
-        Self { cache, storage }
+        Self {
+            cache,
+            storage,
+            prefetched_roots: Arc::new(parking_lot::Mutex::new(HashMap::new())),
+        }
     }
 
     pub async fn get<'new, K: ArrowReadableKey<'new> + 'new>(
@@ -580,6 +599,41 @@ impl RootManager {
         // This will be replaced with a full prefix-based storage shortly
         format!("sparse_index/{}", id)
     }
+
+    fn should_prefetch(&self, id: &Uuid) -> bool {
+        let mut lock_guard = self.prefetched_roots.lock();
+        let expires_at = lock_guard.get(id);
+        match expires_at {
+            Some(expires_at) => {
+                if SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .expect("Do not deploy before UNIX epoch")
+                    < *expires_at
+                {
+                    false
+                } else {
+                    lock_guard.insert(
+                        *id,
+                        SystemTime::now()
+                            .duration_since(UNIX_EPOCH)
+                            .expect("Do not deploy before UNIX epoch")
+                            + std::time::Duration::from_secs(PREFETCH_TTL_HOURS * 3600),
+                    );
+                    true
+                }
+            }
+            None => {
+                lock_guard.insert(
+                    *id,
+                    SystemTime::now()
+                        .duration_since(UNIX_EPOCH)
+                        .expect("Do not deploy before UNIX epoch")
+                        + std::time::Duration::from_secs(PREFETCH_TTL_HOURS * 3600),
+                );
+                true
+            }
+        }
+    }
 }
 
 #[cfg(test)]

rust/segment/src/blockfile_metadata.rs

@@ -12,6 +12,11 @@ use chroma_index::metadata::types::{
     MetadataIndexError, MetadataIndexFlusher, MetadataIndexReader, MetadataIndexWriter,
 };
 use chroma_types::SegmentType;
+use chroma_types::BOOL_METADATA;
+use chroma_types::F32_METADATA;
+use chroma_types::FULL_TEXT_PLS;
+use chroma_types::STRING_METADATA;
+use chroma_types::U32_METADATA;
 use chroma_types::{MaterializedLogOperation, MetadataValue, Segment, SegmentUuid};
 use core::panic;
 use roaring::RoaringBitmap;
@@ -21,12 +26,6 @@ use tantivy::tokenizer::NgramTokenizer;
 use thiserror::Error;
 use uuid::Uuid;
 
-const FULL_TEXT_PLS: &str = "full_text_pls";
-const STRING_METADATA: &str = "string_metadata";
-const BOOL_METADATA: &str = "bool_metadata";
-const F32_METADATA: &str = "f32_metadata";
-const U32_METADATA: &str = "u32_metadata";
-
 #[derive(Clone)]
 pub struct MetadataSegmentWriter<'me> {
     pub(crate) full_text_index_writer: Option<FullTextIndexWriter>,

rust/segment/src/blockfile_record.rs

@@ -6,7 +6,10 @@ use chroma_blockstore::{
 };
 use chroma_error::{ChromaError, ErrorCodes};
 use chroma_index::fulltext::types::FullTextIndexError;
-use chroma_types::{DataRecord, MaterializedLogOperation, Segment, SegmentType, SegmentUuid};
+use chroma_types::{
+    DataRecord, MaterializedLogOperation, Segment, SegmentType, SegmentUuid, MAX_OFFSET_ID,
+    OFFSET_ID_TO_DATA, OFFSET_ID_TO_USER_ID, USER_ID_TO_OFFSET_ID,
+};
 use futures::{Stream, StreamExt, TryStreamExt};
 use std::collections::HashMap;
 use std::fmt::{self, Debug, Formatter};
@@ -16,11 +19,6 @@ use std::sync::Arc;
 use thiserror::Error;
 use uuid::Uuid;
 
-const USER_ID_TO_OFFSET_ID: &str = "user_id_to_offset_id";
-const OFFSET_ID_TO_USER_ID: &str = "offset_id_to_user_id";
-const OFFSET_ID_TO_DATA: &str = "offset_id_to_data";
-const MAX_OFFSET_ID: &str = "max_offset_id";
-
 #[derive(Clone)]
 pub struct RecordSegmentWriter {
     // These are Option<> so that we can take() them when we commit

rust/segment/src/distributed_spann.rs

@@ -18,18 +18,17 @@ use chroma_index::{hnsw_provider::HnswIndexProvider, spann::types::SpannIndexWri
 use chroma_types::Collection;
 use chroma_types::InternalSpannConfiguration;
 use chroma_types::SegmentUuid;
+use chroma_types::HNSW_PATH;
+use chroma_types::MAX_HEAD_ID_BF_PATH;
+use chroma_types::POSTING_LIST_PATH;
+use chroma_types::VERSION_MAP_PATH;
 use chroma_types::{MaterializedLogOperation, Segment, SegmentScope, SegmentType};
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::fmt::Formatter;
 use thiserror::Error;
 use uuid::Uuid;
 
-const HNSW_PATH: &str = "hnsw_path";
-const VERSION_MAP_PATH: &str = "version_map_path";
-const POSTING_LIST_PATH: &str = "posting_list_path";
-const MAX_HEAD_ID_BF_PATH: &str = "max_head_id_path";
-
 #[derive(Clone)]
 pub struct SpannSegmentWriter {
     index: SpannIndexWriter,

rust/types/src/segment.rs

@@ -9,6 +9,22 @@ use thiserror::Error;
 use tonic::Status;
 use uuid::Uuid;
 
+pub const USER_ID_TO_OFFSET_ID: &str = "user_id_to_offset_id";
+pub const OFFSET_ID_TO_USER_ID: &str = "offset_id_to_user_id";
+pub const OFFSET_ID_TO_DATA: &str = "offset_id_to_data";
+pub const MAX_OFFSET_ID: &str = "max_offset_id";
+
+pub const FULL_TEXT_PLS: &str = "full_text_pls";
+pub const STRING_METADATA: &str = "string_metadata";
+pub const BOOL_METADATA: &str = "bool_metadata";
+pub const F32_METADATA: &str = "f32_metadata";
+pub const U32_METADATA: &str = "u32_metadata";
+
+pub const HNSW_PATH: &str = "hnsw_path";
+pub const VERSION_MAP_PATH: &str = "version_map_path";
+pub const POSTING_LIST_PATH: &str = "posting_list_path";
+pub const MAX_HEAD_ID_BF_PATH: &str = "max_head_id_path";
+
 /// SegmentUuid is a wrapper around Uuid to provide a type for the segment id.
 #[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
 pub struct SegmentUuid(pub Uuid);
@@ -94,6 +110,33 @@ pub struct Segment {
     pub file_path: HashMap<String, Vec<String>>,
 }
 
+impl Segment {
+    pub fn prefetch_supported(&self) -> bool {
+        matches!(
+            self.r#type,
+            SegmentType::BlockfileMetadata | SegmentType::BlockfileRecord | SegmentType::Spann
+        )
+    }
+
+    pub fn filepaths_to_prefetch(&self) -> Vec<String> {
+        let mut res = Vec::new();
+        match self.r#type {
+            SegmentType::Spann => {
+                if let Some(pl_path) = self.file_path.get(POSTING_LIST_PATH) {
+                    res.extend(pl_path.iter().cloned());
+                }
+            }
+            SegmentType::BlockfileMetadata | SegmentType::BlockfileRecord => {
+                for paths in self.file_path.values() {
+                    res.extend(paths.iter().cloned());
+                }
+            }
+            _ => {}
+        }
+        res
+    }
+}
+
 #[derive(Error, Debug)]
 pub enum SegmentConversionError {
     #[error("Invalid UUID")]

rust/worker/src/execution/operators/prefetch_segment.rs

@@ -67,27 +67,18 @@ impl Operator<PrefetchSegmentInput, PrefetchSegmentOutput> for PrefetchSegmentOp
         &self,
         input: &PrefetchSegmentInput,
     ) -> Result<PrefetchSegmentOutput, PrefetchSegmentError> {
-        if input.segment.r#type != SegmentType::BlockfileMetadata
-            && input.segment.r#type != SegmentType::BlockfileRecord
-        {
+        if !input.segment.prefetch_supported() {
             return Err(PrefetchSegmentError::UnsupportedSegmentType(
                 input.segment.r#type,
             ));
         }
 
-        tracing::info!(
-            "Prefetching segment: {:?} ({:?})",
-            input.segment.r#type,
-            input.segment.id,
-        );
-
         let mut futures = input
             .segment
-            .file_path
-            .values()
-            .flatten()
+            .filepaths_to_prefetch()
+            .into_iter()
             .map(|blockfile_id| async move {
-                let blockfile_id = Uuid::parse_str(blockfile_id)?;
+                let blockfile_id = Uuid::parse_str(&blockfile_id)?;
                 let count = input.blockfile_provider.prefetch(&blockfile_id).await?;
                 Ok::<_, PrefetchSegmentError>(count)
             })

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

@@ -644,14 +644,14 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             Some(writer) => writer,
             None => return,
         };
-        let vector_writer = match vector_segment.r#type {
+        let (vector_writer, is_vector_segment_spann) = match vector_segment.r#type {
             SegmentType::Spann => match self.ok_or_terminate(
                 self.spann_provider
                     .write(&collection, &vector_segment, dimension)
                     .await,
                 ctx,
             ) {
-                Some(writer) => VectorSegmentWriter::Spann(writer),
+                Some(writer) => (VectorSegmentWriter::Spann(writer), true),
                 None => return,
             },
             _ => match self.ok_or_terminate(
@@ -665,7 +665,7 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
                 .map_err(|err| *err),
                 ctx,
             ) {
-                Some(writer) => VectorSegmentWriter::Hnsw(writer),
+                Some(writer) => (VectorSegmentWriter::Hnsw(writer), false),
                 None => return,
             },
         };
@@ -690,7 +690,13 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
         // Prefetch segments
         let prefetch_segments = match self.rebuild {
             true => vec![output.record_segment],
-            false => vec![output.metadata_segment, output.record_segment],
+            false => {
+                let mut segments = vec![output.metadata_segment, output.record_segment];
+                if is_vector_segment_spann {
+                    segments.push(output.vector_segment);
+                }
+                segments
+            }
         };
         for segment in prefetch_segments {
             let prefetch_task = wrap(

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

@@ -17,6 +17,9 @@ use crate::execution::operators::{
     prefetch_record::{
         PrefetchRecordError, PrefetchRecordInput, PrefetchRecordOperator, PrefetchRecordOutput,
     },
+    prefetch_segment::{
+        PrefetchSegmentError, PrefetchSegmentInput, PrefetchSegmentOperator, PrefetchSegmentOutput,
+    },
     spann_bf_pl::{SpannBfPlError, SpannBfPlInput, SpannBfPlOperator, SpannBfPlOutput},
     spann_centers_search::{
         SpannCentersSearchError, SpannCentersSearchInput, SpannCentersSearchOperator,
@@ -170,6 +173,16 @@ impl Orchestrator for SpannKnnOrchestrator {
         );
         tasks.push(head_search_task);
 
+        let prefetch_task = wrap(
+            Box::new(PrefetchSegmentOperator::new()),
+            PrefetchSegmentInput::new(
+                self.knn_filter_output.vector_segment.clone(),
+                self.spann_provider.blockfile_provider.clone(),
+            ),
+            ctx.receiver(),
+        );
+        tasks.push(prefetch_task);
+
         tasks
     }
 
@@ -188,6 +201,19 @@ impl Orchestrator for SpannKnnOrchestrator {
     }
 }
 
+#[async_trait]
+impl Handler<TaskResult<PrefetchSegmentOutput, PrefetchSegmentError>> for SpannKnnOrchestrator {
+    type Result = ();
+
+    async fn handle(
+        &mut self,
+        _: TaskResult<PrefetchSegmentOutput, PrefetchSegmentError>,
+        _: &ComponentContext<SpannKnnOrchestrator>,
+    ) {
+        // Nothing to do.
+    }
+}
+
 #[async_trait]
 impl Handler<TaskResult<KnnLogOutput, KnnLogError>> for SpannKnnOrchestrator {
     type Result = ();