[ENH] Load blocks given keys under P0 priority (#6006)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
- Rename `prefetch_*` to `load_*` to better indicate its usage. Prefetch
is for lower priority data warmup, and should be used in a fire and
forget fashion. Load higher priority and could be blocked on.
- Make `load_blocks` use `P0` priority instead of `P1` when fetching
blocks. Previously prefetch use `P1` priority and led to long tail
latency on log materialization.
- New functionality
  - N/A

## Test plan

_How are these changes tested?_

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility
changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: Sicheng-Pan

rust/blockstore/src/arrow/blockfile.rs

@@ -526,7 +526,7 @@ impl<'me, K: ArrowReadableKey<'me> + Into<KeyWrapper>, V: ArrowReadableValue<'me
             if !self.block_manager.cached(block_id).await
                 && !self.loaded_blocks.read().contains_key(block_id)
             {
-                futures.push(self.get_block(*block_id, StorageRequestPriority::P1));
+                futures.push(self.get_block(*block_id, StorageRequestPriority::P0));
             }
         }
         join_all(futures).await;

rust/blockstore/src/types/reader.rs

@@ -137,9 +137,10 @@ impl<
         }
     }
 
-    pub async fn load_blocks_for_keys(&self, keys: impl IntoIterator<Item = (String, K)>) {
+    // NOTE(sicheng): This loads the underlying data concurrently
+    pub async fn load_data_for_keys(&self, keys: impl IntoIterator<Item = (String, K)>) {
         match self {
-            BlockfileReader::MemoryBlockfileReader(_reader) => unimplemented!(),
+            BlockfileReader::MemoryBlockfileReader(_) => (),
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader.load_blocks_for_keys(keys).await
             }

rust/index/src/fulltext/types.rs

@@ -440,7 +440,7 @@ impl NgramLiteralProvider<FullTextIndexError> for FullTextIndexReader<'_> {
         6
     }
 
-    async fn prefetch_ngrams<'me, Ngrams>(&'me self, ngrams: Ngrams)
+    async fn load_ngrams<'me, Ngrams>(&'me self, ngrams: Ngrams)
     where
         Ngrams: IntoIterator<Item = &'me str> + Send + Sync,
     {

rust/segment/src/blockfile_record.rs

@@ -913,15 +913,15 @@ impl RecordSegmentReader<'_> {
         self.id_to_user_id.count().await
     }
 
-    pub async fn prefetch_id_to_data(&self, keys: &[u32]) {
+    pub async fn load_id_to_data(&self, keys: impl Iterator<Item = u32>) {
         self.id_to_data
-            .load_blocks_for_keys(keys.iter().map(|k| ("".to_string(), *k)))
+            .load_data_for_keys(keys.map(|k| ("".to_string(), k)))
             .await
     }
 
-    pub(crate) async fn prefetch_user_id_to_id(&self, keys: &[&str]) {
+    pub async fn load_user_id_to_id(&self, keys: impl Iterator<Item = &str>) {
         self.user_id_to_id
-            .load_blocks_for_keys(keys.iter().map(|k| ("".to_string(), *k)))
+            .load_data_for_keys(keys.map(|k| ("".to_string(), k)))
             .await
     }
 

rust/segment/src/types.rs

@@ -597,7 +597,7 @@ pub async fn materialize_logs(
         user_ids.sort_unstable();
         user_ids.dedup();
         async {
-            reader.prefetch_user_id_to_id(&user_ids).await;
+            reader.load_user_id_to_id(user_ids.iter().cloned()).await;
 
             let mut existing_offset_ids = Vec::with_capacity(user_ids.len());
             for user_id in user_ids {
@@ -610,7 +610,9 @@ pub async fn materialize_logs(
                 };
             }
 
-            reader.prefetch_id_to_data(&existing_offset_ids).await;
+            reader
+                .load_id_to_data(existing_offset_ids.iter().cloned())
+                .await;
             Ok::<_, LogMaterializerError>(())
         }
         .instrument(Span::current())

rust/types/src/regex/literal_expr.rs

@@ -79,7 +79,7 @@ pub trait NgramLiteralProvider<E, const N: usize = 3> {
     // Return the max branching factor during the search
     fn maximum_branching_factor(&self) -> usize;
 
-    async fn prefetch_ngrams<'me, Ngrams>(&'me self, _ngrams: Ngrams)
+    async fn load_ngrams<'me, Ngrams>(&'me self, _ngrams: Ngrams)
     where
         Ngrams: IntoIterator<Item = &'me str> + Send + Sync,
     {
@@ -180,7 +180,7 @@ pub trait NgramLiteralProvider<E, const N: usize = 3> {
             return Ok(HashSet::new());
         }
 
-        self.prefetch_ngrams(
+        self.load_ngrams(
             ngram_vec
                 .iter()
                 .flat_map(|ngrams| ngrams.iter().map(|ngram| ngram.as_str())),