[ENH] Add orchestration for count (#2180)

## Description of changes

*Summarize the changes made by this PR.*
 - New functionality
Introduces the orchestrator for count and a basic counting logic where
we simply add the number of log records with the number of records in
the blockfile. More advanced merging like deduping adds and updates and
removing deletes will be a different PR.

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

## Documentation Changes
None

---------

Co-authored-by: Sanket Kedia <[email protected]>

Author: sanketkedia

rust/worker/src/blockstore/arrow/blockfile.rs

@@ -277,11 +277,17 @@ impl<'me, K: ArrowReadableKey<'me>, V: ArrowReadableValue<'me>> ArrowBlockfileRe
 
     // Count the total number of records.
     pub(crate) async fn count(&self) -> Result<usize, Box<dyn ChromaError>> {
+        let mut block_ids: Vec<Uuid> = vec![];
+        {
+            let lock_guard = self.sparse_index.forward.lock();
+            let mut curr_iter = lock_guard.iter();
+            while let Some((_, block_id)) = curr_iter.next() {
+                block_ids.push(block_id.clone());
+            }
+        }
         let mut result: usize = 0;
-        let lock_guard = self.sparse_index.forward.lock();
-        let mut curr_iter = lock_guard.iter();
-        while let Some((_, block_id)) = curr_iter.next() {
-            let block = self.get_block(*block_id).await;
+        for block_id in block_ids {
+            let block = self.get_block(block_id).await;
             match block {
                 Some(b) => result = result + b.len(),
                 None => {

rust/worker/src/blockstore/types.rs

@@ -27,6 +27,8 @@ pub(crate) enum BlockfileError {
     TransactionInProgress,
     #[error("Transaction not in progress")]
     TransactionNotInProgress,
+    #[error("Block not found")]
+    BlockNotFound,
 }
 
 impl ChromaError for BlockfileError {
@@ -38,6 +40,7 @@ impl ChromaError for BlockfileError {
             BlockfileError::TransactionInProgress | BlockfileError::TransactionNotInProgress => {
                 ErrorCodes::FailedPrecondition
             }
+            BlockfileError::BlockNotFound => ErrorCodes::Internal,
         }
     }
 }
@@ -297,7 +300,17 @@ impl<
     pub(crate) async fn count(&'referred_data self) -> Result<usize, Box<dyn ChromaError>> {
         match self {
             BlockfileReader::MemoryBlockfileReader(reader) => reader.count(),
-            BlockfileReader::ArrowBlockfileReader(reader) => reader.count().await,
+            BlockfileReader::ArrowBlockfileReader(reader) => {
+                let count = reader.count().await;
+                match count {
+                    Ok(c) => {
+                        return Ok(c);
+                    }
+                    Err(_) => {
+                        return Err(Box::new(BlockfileError::BlockNotFound));
+                    }
+                }
+            }
         }
     }
 

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

@@ -0,0 +1,90 @@
+use thiserror::Error;
+
+use tonic::async_trait;
+
+use crate::{
+    blockstore::provider::BlockfileProvider,
+    errors::{ChromaError, ErrorCodes},
+    execution::operator::Operator,
+    segment::record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
+    types::Segment,
+};
+
+#[derive(Debug)]
+pub(crate) struct CountRecordsOperator {}
+
+impl CountRecordsOperator {
+    pub(crate) fn new() -> Box<Self> {
+        Box::new(CountRecordsOperator {})
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct CountRecordsInput {
+    record_segment_definition: Segment,
+    blockfile_provider: BlockfileProvider,
+}
+
+impl CountRecordsInput {
+    pub(crate) fn new(
+        record_segment_definition: Segment,
+        blockfile_provider: BlockfileProvider,
+    ) -> Self {
+        Self {
+            record_segment_definition,
+            blockfile_provider,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct CountRecordsOutput {
+    pub(crate) count: usize,
+}
+
+#[derive(Error, Debug)]
+pub(crate) enum CountRecordsError {
+    #[error("Error reading record segment reader")]
+    RecordSegmentReadError,
+    #[error("Error creating record segment reader")]
+    RecordSegmentError(#[from] RecordSegmentReaderCreationError),
+}
+
+impl ChromaError for CountRecordsError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            CountRecordsError::RecordSegmentError(_) => ErrorCodes::Internal,
+            CountRecordsError::RecordSegmentReadError => ErrorCodes::Internal,
+        }
+    }
+}
+
+#[async_trait]
+impl Operator<CountRecordsInput, CountRecordsOutput> for CountRecordsOperator {
+    type Error = CountRecordsError;
+    async fn run(
+        &self,
+        input: &CountRecordsInput,
+    ) -> Result<CountRecordsOutput, CountRecordsError> {
+        let segment_reader = RecordSegmentReader::from_segment(
+            &input.record_segment_definition,
+            &input.blockfile_provider,
+        )
+        .await;
+        match segment_reader {
+            Ok(reader) => match reader.count().await {
+                Ok(val) => {
+                    return Ok(CountRecordsOutput { count: val });
+                }
+                Err(_) => {
+                    println!("Error reading record segment");
+                    return Err(CountRecordsError::RecordSegmentReadError);
+                }
+            },
+            Err(e) => {
+                println!("Error opening record segment");
+                return Err(CountRecordsError::RecordSegmentError(*e));
+            }
+        }
+    }
+}

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

@@ -1,4 +1,5 @@
 pub(super) mod brute_force_knn;
+pub(super) mod count_records;
 pub(super) mod flush_s3;
 pub(super) mod hnsw_knn;
 pub(super) mod merge_knn_results;