Cleanup

Author: codetheweb

rust/blockstore/src/arrow/blockfile.rs

@@ -11,11 +11,9 @@ use crate::arrow::root::CURRENT_VERSION;
 use crate::arrow::sparse_index::SparseIndexWriter;
 use crate::key::CompositeKey;
 use crate::key::KeyWrapper;
-use crate::Key;
 use chroma_error::ChromaError;
 use chroma_error::ErrorCodes;
 use futures::future::join_all;
-use futures::stream::FuturesUnordered;
 use futures::{Stream, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use std::collections::HashSet;
@@ -344,42 +342,6 @@ impl ArrowUnorderedBlockfileWriter {
         Ok(())
     }
 
-    pub(crate) async fn prefetch<K: Key>(
-        &self,
-        keys: impl IntoIterator<Item = (String, K)>,
-    ) -> Result<(), Box<dyn ChromaError>> {
-        let mut block_ids = HashSet::new();
-
-        for key in keys {
-            let block_id = self
-                .root
-                .sparse_index
-                .get_target_block_id(&CompositeKey::new(key.0, key.1));
-            block_ids.insert(block_id);
-        }
-
-        let mut block_fetches = FuturesUnordered::new();
-        for block_id in &block_ids {
-            block_fetches.push(self.block_manager.get(block_id));
-        }
-
-        while let Some(block) = block_fetches.next().await {
-            match block {
-                Ok(_) => {}
-                Err(err) => match err {
-                    GetError::StorageGetError(chroma_storage::GetError::NoSuchKey(_)) => {
-                        // Ignore
-                    }
-                    _ => {
-                        return Err(Box::new(err) as Box<dyn ChromaError>);
-                    }
-                },
-            }
-        }
-
-        Ok(())
-    }
-
     pub(crate) fn id(&self) -> Uuid {
         self.id
     }

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

@@ -10,15 +10,11 @@ use super::{
     flusher::ArrowBlockfileFlusher,
     types::{ArrowWriteableKey, ArrowWriteableValue},
 };
-use crate::arrow::provider::GetError;
 use crate::arrow::root::CURRENT_VERSION;
 use crate::arrow::sparse_index::SparseIndexWriter;
 use crate::key::CompositeKey;
-use crate::Key;
 use chroma_error::ChromaError;
 use chroma_error::ErrorCodes;
-use futures::stream::FuturesUnordered;
-use futures::StreamExt;
 use itertools::Itertools;
 use std::collections::HashSet;
 use std::collections::VecDeque;
@@ -156,11 +152,16 @@ impl ArrowOrderedBlockfileWriter {
             .load(std::sync::atomic::Ordering::Relaxed);
         let block_cache_hit_rate = 100.0 * num_block_cache_hits as f64
             / (num_block_cache_hits + num_block_cache_misses) as f64;
-        tracing::debug!(
-            block_cache_hit_rate = block_cache_hit_rate,
-            "Block cache hit rate: {:.2}%",
-            block_cache_hit_rate,
-        );
+        if block_cache_hit_rate.is_nan() {
+            // This can happen if there are no cache hits or misses
+            tracing::debug!("Block cache hit rate: N/A");
+        } else {
+            tracing::debug!(
+                block_cache_hit_rate = block_cache_hit_rate,
+                "Block cache hit rate: {:.2}%",
+                block_cache_hit_rate,
+            );
+        }
 
         let mut split_block_deltas = Vec::new();
         for delta in inner.completed_block_deltas.drain(..) {

rust/index/src/fulltext/types.rs

@@ -166,7 +166,10 @@ impl FullTextIndexWriter {
         &self,
         blockfile_provider: &BlockfileProvider,
     ) -> Result<(), FullTextIndexError> {
-        const PERCENT_TO_PREFETCH: f32 = 0.001;
+        // `token_instances` can be millions of items. Imagine there are 10k documents, each with 1000 tokens. There are then ~10000*1000*3=30M token instances.
+        // In comparison, the number of blocks required to store these is much, much smaller. Say (worst case) each token instance consumes 12 bytes of a block (trigram, offset ID, position) and a block is 8MB. Then a single block can store ~660k token instances.
+        // Given this, and assuming that token instances are relatively uniformly distributed across blocks, we can prefetch a very small percentage of the key space and the cache hit rate will still be at or near 100%.
+        const PERCENT_TO_PREFETCH: f32 = 0.0005;
 
         if let Some(id) = self.posting_lists_blockfile_writer.forked_from_id() {
             let mut keys = vec![];

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

@@ -818,7 +818,11 @@ impl Handler<TaskResult<ApplyLogToSegmentWriterOutput, ApplyLogToSegmentWriterOp
         };
 
         if num_tasks_left == 0 {
-            let writers = self.get_segment_writers().await.unwrap();
+            let writers = self.get_segment_writers().await;
+            let writers = match self.ok_or_terminate(writers, ctx) {
+                Some(writers) => writers,
+                None => return,
+            };
 
             if message.segment_id == writers.metadata.id {
                 let span = self.get_segment_writer_span(&ChromaSegmentWriter::MetadataSegment(
@@ -862,7 +866,11 @@ impl Handler<TaskResult<PrefetchForMetadataWriterOutput, PrefetchForMetadataWrit
             None => return,
         };
 
-        let writers = self.get_segment_writers().await.unwrap();
+        let writers = self.get_segment_writers().await;
+        let writers = match self.ok_or_terminate(writers, ctx) {
+            Some(writers) => writers,
+            None => return,
+        };
 
         self.dispatch_segment_writer_commit(
             ChromaSegmentWriter::MetadataSegment(writers.metadata),