Repository: chroma-core/chroma
Pull request #4127: [ENH] Implement SourceRecordSegment operator
Author: login: Sicheng-Pan
Created: 2025-04-02T00:34:53Z | Merged: 2025-04-04T23:45:25Z

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - Update the compact orchestrator to use record segment instead of log service during rebuild
 - New functionality
   - Implement a `SourceRecordSegmentOperator` that generate materialized logs based on data from record segment
   

## Test plan
*How are these changes tested?*

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*



Diff:
diff --git a/rust/segment/src/blockfile_record.rs b/rust/segment/src/blockfile_record.rs
index a3907929975..cd6bd1b0518 100644
--- a/rust/segment/src/blockfile_record.rs
+++ b/rust/segment/src/blockfile_record.rs
@@ -827,6 +827,15 @@ impl RecordSegmentReader<'_> {
             .map(|vec| vec.into_iter().map(|(_, data)| data).collect())
     }
 
+    pub async fn get_data_stream<'me>(
+        &'me self,
+        offset_range: impl RangeBounds<u32> + Clone + Send + 'me,
+    ) -> impl Stream<Item = Result<DataRecord<'me>, Box<dyn ChromaError>>> + 'me {
+        self.id_to_data
+            .get_range_stream(""..="", offset_range)
+            .map(|res| res.map(|(_, rec)| rec))
+    }
+
     /// Get a stream of offset ids from the smallest to the largest in the given range
     pub fn get_offset_stream<'me>(
         &'me self,
diff --git a/rust/worker/src/execution/operators/mod.rs b/rust/worker/src/execution/operators/mod.rs
index 523014b63f3..876158b7bff 100644
--- a/rust/worker/src/execution/operators/mod.rs
+++ b/rust/worker/src/execution/operators/mod.rs
@@ -22,3 +22,4 @@ pub mod partition_log;
 pub mod prefetch_record;
 pub mod prefetch_segment;
 pub mod projection;
+pub mod source_record_segment;
diff --git a/rust/worker/src/execution/operators/source_record_segment.rs b/rust/worker/src/execution/operators/source_record_segment.rs
new file mode 100644
index 00000000000..ab55c9e0624
--- /dev/null
+++ b/rust/worker/src/execution/operators/source_record_segment.rs
@@ -0,0 +1,129 @@
+use async_trait::async_trait;
+use chroma_error::{ChromaError, ErrorCodes};
+use chroma_segment::blockfile_record::RecordSegmentReader;
+use chroma_system::Operator;
+use chroma_types::{Chunk, LogRecord, OperationRecord};
+use futures::{StreamExt, TryStreamExt};
+use thiserror::Error;
+
+/// The `SourceRecordSegmentOperator` generates logs from the data in the record segment
+///
+/// # Parameters
+/// None
+///
+/// # Inputs
+/// - `record_reader`: The record segment reader, if the collection is initialized
+///
+/// # Outputs
+/// - Chunk of addition logs
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentOperator {}
+
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentInput {
+    pub record_segment_reader: Option<RecordSegmentReader<'static>>,
+}
+
+pub type SourceRecordSegmentOutput = Chunk<LogRecord>;
+
+#[derive(Debug, Error)]
+pub enum SourceRecordSegmentError {
+    #[error("Error reading record segment: {0}")]
+    RecordSegment(#[from] Box<dyn ChromaError>),
+}
+
+impl ChromaError for SourceRecordSegmentError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            SourceRecordSegmentError::RecordSegment(e) => e.code(),
+        }
+    }
+}
+
+#[async_trait]
+impl Operator<SourceRecordSegmentInput, SourceRecordSegmentOutput> for SourceRecordSegmentOperator {
+    type Error = SourceRecordSegmentError;
+    async fn run(
+        &self,
+        input: &SourceRecordSegmentInput,
+    ) -> Result<SourceRecordSegmentOutput, SourceRecordSegmentError> {
+        tracing::trace!("[{}]: {:?}", self.get_name(), input);
+        let logs = match input.record_segment_reader.as_ref() {
+            Some(reader) => {
+                reader
+                    .get_data_stream(..)
+                    .await
+                    .enumerate()
+                    .map(|(offset, res)| {
+                        res.map(|rec| LogRecord {
+                            // Log offset starts with 1
+                            log_offset: offset as i64 + 1,
+                            record: OperationRecord {
+                                id: rec.id.to_string(),
+                                embedding: Some(rec.embedding.to_vec()),
+                                encoding: Some(chroma_types::ScalarEncoding::FLOAT32),
+                                metadata: rec.metadata.map(|meta| {
+                                    meta.into_iter().map(|(k, v)| (k, v.into())).collect()
+                                }),
+                                document: rec.document.map(ToString::to_string),
+                                operation: chroma_types::Operation::Add,
+                            },
+                        })
+                    })
+                    .try_collect::<Vec<_>>()
+                    .await?
+            }
+            None => Default::default(),
+        };
+        Ok(Chunk::new(logs.into()))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use chroma_log::test::{int_as_id, upsert_generator, LoadFromGenerator};
+    use chroma_segment::{blockfile_record::RecordSegmentReader, test::TestDistributedSegment};
+    use chroma_system::Operator;
+    use chroma_types::Operation;
+
+    use crate::execution::operators::source_record_segment::SourceRecordSegmentOperator;
+
+    use super::SourceRecordSegmentInput;
+
+    /// The unit tests for `SourceRecordSegmentOperator` uses the following test data
+    /// It generates 100 log records and compact them
+    async fn setup_source_input() -> SourceRecordSegmentInput {
+        let mut test_segment = TestDistributedSegment::default();
+        test_segment
+            .populate_with_generator(100, upsert_generator)
+            .await;
+        let reader = RecordSegmentReader::from_segment(
+            &test_segment.record_segment,
+            &test_segment.blockfile_provider,
+        )
+        .await
+        .expect("Record segment reader should be initialized");
+        SourceRecordSegmentInput {
+            record_segment_reader: Some(reader),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_source() {
+        let source_input = setup_source_input().await;
+
+        let source_operator = SourceRecordSegmentOperator {};
+
+        let source_output = source_operator
+            .run(&source_input)
+            .await
+            .expect("SourceOperator should not fail");
+
+        assert_eq!(source_output.len(), 100);
+        for (offset, (record, _)) in source_output.iter().enumerate() {
+            assert_eq!(record.log_offset, offset as i64 + 1);
+            assert_eq!(record.record.id, int_as_id(offset + 1));
+            assert_eq!(record.record.operation, Operation::Add);
+        }
+    }
+}
diff --git a/rust/worker/src/execution/orchestration/compact.rs b/rust/worker/src/execution/orchestration/compact.rs
index d6fd2ce30d3..99a4baf17b2 100644
--- a/rust/worker/src/execution/orchestration/compact.rs
+++ b/rust/worker/src/execution/orchestration/compact.rs
@@ -25,7 +25,7 @@ use chroma_segment::{
 use chroma_sysdb::SysDb;
 use chroma_system::{
     wrap, ChannelError, ComponentContext, ComponentHandle, Dispatcher, Handler, Orchestrator,
-    PanicError, ReceiverForMessage, TaskError, TaskMessage, TaskResult,
+    PanicError, TaskError, TaskMessage, TaskResult,
 };
 use chroma_types::{Chunk, LogRecord, SegmentFlushInfo, SegmentType, SegmentUuid};
 use thiserror::Error;
@@ -62,6 +62,10 @@ use crate::{
             PrefetchSegmentOutput,
         },
         register::{RegisterError, RegisterInput, RegisterOperator, RegisterOutput},
+        source_record_segment::{
+            SourceRecordSegmentError, SourceRecordSegmentInput, SourceRecordSegmentOperator,
+            SourceRecordSegmentOutput,
+        },
     },
 };
 
@@ -109,8 +113,6 @@ pub struct CompactOrchestrator {
     blockfile_provider: BlockfileProvider,
     hnsw_provider: HnswIndexProvider,
     spann_provider: SpannProvider,
-    // State we hold across the execution
-    pulled_log_offset: Option<i64>,
     // Dispatcher
     dispatcher: ComponentHandle<Dispatcher>,
     // Tracks the total remaining number of MaterializeLogs tasks
@@ -119,6 +121,8 @@ pub struct CompactOrchestrator {
     num_uncompleted_tasks_by_segment: HashMap<SegmentUuid, usize>,
     // Tracks the total collection size in number of bytes
     collection_logical_size_bytes: i64,
+    // Tracks the last log offset we have witnessed
+    pulled_log_offset: i64,
     // Result Channel
     result_channel: Option<Sender<Result<CompactionResponse, CompactionError>>>,
     max_compaction_size: usize,
@@ -174,6 +178,8 @@ pub enum CompactionError {
     Result(#[from] RecvError),
     #[error("Error creaitng spann writer: {0}")]
     SpannSegment(#[from] SpannSegmentWriterError),
+    #[error("Error sourcing record segment: {0}")]
+    SourceRecordSegment(#[from] SourceRecordSegmentError),
 }
 
 impl<E> From<TaskError<E>> for CompactionError
@@ -226,7 +232,7 @@ impl CompactOrchestrator {
             blockfile_provider,
             hnsw_provider,
             spann_provider,
-            pulled_log_offset: None,
+            pulled_log_offset: 0,
             dispatcher,
             num_uncompleted_materialization_tasks: 0,
             num_uncompleted_tasks_by_segment: HashMap::new(),
@@ -291,11 +297,6 @@ impl CompactOrchestrator {
     async fn dispatch_apply_log_to_segment_writer_tasks(
         &mut self,
         materialized_logs: MaterializeLogsResult,
-        self_address: Box<
-            dyn ReceiverForMessage<
-                TaskResult<ApplyLogToSegmentWriterOutput, ApplyLogToSegmentWriterOperatorError>,
-            >,
-        >,
         ctx: &ComponentContext<CompactOrchestrator>,
     ) {
         let writers = match self.ok_or_terminate(self.get_segment_writers(), ctx) {
@@ -303,15 +304,15 @@ impl CompactOrchestrator {
             None => return,
         };
 
-        {
+        if !self.compaction_job.rebuild {
             self.num_uncompleted_tasks_by_segment
-                .entry(writers.metadata_writer.id)
+                .entry(writers.record_writer.id)
                 .and_modify(|v| {
                     *v += 1;
                 })
                 .or_insert(1);
 
-            let writer = ChromaSegmentWriter::MetadataSegment(writers.metadata_writer);
+            let writer = ChromaSegmentWriter::RecordSegment(writers.record_writer);
             let span = self.get_segment_writer_span(&writer);
             let operator = ApplyLogToSegmentWriterOperator::new();
             let input = ApplyLogToSegmentWriterInput::new(
@@ -319,23 +320,22 @@ impl CompactOrchestrator {
                 materialized_logs.clone(),
                 writers.record_reader.clone(),
             );
-            let task = wrap(operator, input, self_address.clone());
+            let task = wrap(operator, input, ctx.receiver());
             let res = self.dispatcher().send(task, Some(span)).await;
-            match self.ok_or_terminate(res, ctx) {
-                Some(_) => (),
-                None => return,
+            if self.ok_or_terminate(res, ctx).is_none() {
+                return;
             }
         }
 
         {
             self.num_uncompleted_tasks_by_segment
-                .entry(writers.record_writer.id)
+                .entry(writers.metadata_writer.id)
                 .and_modify(|v| {
                     *v += 1;
                 })
                 .or_insert(1);
 
-            let writer = ChromaSegmentWriter::RecordSegment(writers.record_writer);
+            let writer = ChromaSegmentWriter::MetadataSegment(writers.metadata_writer);
             let span = self.get_segment_writer_span(&writer);
             let operator = ApplyLogToSegmentWriterOperator::new();
             let input = ApplyLogToSegmentWriterInput::new(
@@ -343,11 +343,10 @@ impl CompactOrchestrator {
                 materialized_logs.clone(),
                 writers.record_reader.clone(),
             );
-            let task = wrap(operator, input, self_address.clone());
+            let task = wrap(operator, input, ctx.receiver());
             let res = self.dispatcher().send(task, Some(span)).await;
-            match self.ok_or_terminate(res, ctx) {
-                Some(_) => (),
-                None => return,
+            if self.ok_or_terminate(res, ctx).is_none() {
+                return;
             }
         }
 
@@ -364,7 +363,7 @@ impl CompactOrchestrator {
             let operator = ApplyLogToSegmentWriterOperator::new();
             let input =
                 ApplyLogToSegmentWriterInput::new(writer, materialized_logs, writers.record_reader);
-            let task = wrap(operator, input, self_address);
+            let task = wrap(operator, input, ctx.receiver());
             let res = self.dispatcher().send(task, Some(span)).await;
             self.ok_or_terminate(res, ctx);
         }
@@ -373,17 +372,12 @@ impl CompactOrchestrator {
     async fn dispatch_segment_writer_commit(
         &mut self,
         segment_writer: ChromaSegmentWriter<'static>,
-        self_address: Box<
-            dyn ReceiverForMessage<
-                TaskResult<CommitSegmentWriterOutput, CommitSegmentWriterOperatorError>,
-            >,
-        >,
         ctx: &ComponentContext<CompactOrchestrator>,
     ) {
         let span = self.get_segment_writer_span(&segment_writer);
         let operator = CommitSegmentWriterOperator::new();
         let input = CommitSegmentWriterInput::new(segment_writer);
-        let task = wrap(operator, input, self_address);
+        let task = wrap(operator, input, ctx.receiver());
         let res = self.dispatcher().send(task, Some(span)).await;
         self.ok_or_terminate(res, ctx);
     }
@@ -391,28 +385,23 @@ impl CompactOrchestrator {
     async fn dispatch_segment_flush(
         &mut self,
         segment_flusher: ChromaSegmentFlusher,
-        self_address: Box<
-            dyn ReceiverForMessage<
-                TaskResult<FlushSegmentWriterOutput, FlushSegmentWriterOperatorError>,
-            >,
-        >,
         ctx: &ComponentContext<CompactOrchestrator>,
     ) {
         let span = self.get_segment_flusher_span(&segment_flusher);
         let operator = FlushSegmentWriterOperator::new();
         let input = FlushSegmentWriterInput::new(segment_flusher);
-        let task = wrap(operator, input, self_address);
+        let task = wrap(operator, input, ctx.receiver());
         let res = self.dispatcher().send(task, Some(span)).await;
         self.ok_or_terminate(res, ctx);
     }
 
-    async fn register(&mut self, log_position: i64, ctx: &ComponentContext<CompactOrchestrator>) {
+    async fn register(&mut self, ctx: &ComponentContext<CompactOrchestrator>) {
         self.state = ExecutionState::Register;
         let operator = RegisterOperator::new();
         let input = RegisterInput::new(
             self.compaction_job.tenant_id.clone(),
             self.compaction_job.collection_id,
-            log_position,
+            self.pulled_log_offset,
             self.compaction_job.collection_version,
             self.flush_results.clone().into(),
             self.total_records_last_compaction,
@@ -557,6 +546,7 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             true => 0,
             false => collection.size_bytes_post_compaction as i64,
         };
+        self.pulled_log_offset = collection.log_position;
 
         let mut metadata_segment = output.metadata_segment.clone();
         let mut vector_segment = output.vector_segment.clone();
@@ -583,8 +573,7 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
         ) {
             Some(reader) => reader,
             None => return,
-        }
-        .filter(|_| !self.compaction_job.rebuild);
+        };
         let record_writer = match self.ok_or_terminate(
             RecordSegmentWriter::from_segment(&output.record_segment, &self.blockfile_provider)
                 .await,
@@ -627,7 +616,9 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
         };
 
         let writers = CompactWriters {
-            record_reader,
+            record_reader: record_reader
+                .clone()
+                .filter(|_| !self.compaction_job.rebuild),
             metadata_writer,
             record_writer,
             vector_writer,
@@ -643,32 +634,44 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             return;
         }
 
-        // Prefetch metadata and record segment
-        let prefetch_tasks = [output.metadata_segment, output.record_segment].map(|segment| {
-            wrap(
+        // Prefetch segments
+        let prefetch_segments = match self.compaction_job.rebuild {
+            true => vec![output.record_segment],
+            false => vec![output.metadata_segment, output.record_segment],
+        };
+        for segment in prefetch_segments {
+            let prefetch_task = wrap(
                 Box::new(PrefetchSegmentOperator::new()),
                 PrefetchSegmentInput::new(segment, self.blockfile_provider.clone()),
                 ctx.receiver(),
-            )
-        });
-        for task in prefetch_tasks {
-            self.send(task, ctx).await;
+            );
+            self.send(prefetch_task, ctx).await;
         }
 
-        let fetch_log_task = wrap(
-            Box::new(FetchLogOperator {
-                log_client: self.log.clone(),
-                batch_size: self.fetch_log_batch_size,
-                // Here we do not need to be inclusive since the compaction job
-                // offset is the one after the last compaction offset
-                start_log_offset_id: self.compaction_job.offset as u32,
-                maximum_fetch_count: Some(self.max_compaction_size as u32),
-                collection_uuid: self.compaction_job.collection_id,
-            }),
-            (),
-            ctx.receiver(),
-        );
-        self.send(fetch_log_task, ctx).await;
+        let log_task = match self.compaction_job.rebuild {
+            true => wrap(
+                Box::new(SourceRecordSegmentOperator {}),
+                SourceRecordSegmentInput {
+                    record_segment_reader: record_reader,
+                },
+                ctx.receiver(),
+            ),
+            false => wrap(
+                Box::new(FetchLogOperator {
+                    log_client: self.log.clone(),
+                    batch_size: self.fetch_log_batch_size,
+                    // Here we do not need to be inclusive since the compaction job
+                    // offset is the one after the last compaction offset
+                    start_log_offset_id: self.compaction_job.offset as u32,
+                    maximum_fetch_count: Some(self.max_compaction_size as u32),
+                    collection_uuid: self.compaction_job.collection_id,
+                }),
+                (),
+                ctx.receiver(),
+            ),
+        };
+
+        self.send(log_task, ctx).await;
     }
 }
 
@@ -694,34 +697,50 @@ impl Handler<TaskResult<FetchLogOutput, FetchLogError>> for CompactOrchestrator
         message: TaskResult<FetchLogOutput, FetchLogError>,
         ctx: &ComponentContext<CompactOrchestrator>,
     ) {
-        let records = match self.ok_or_terminate(message.into_inner(), ctx) {
+        let output = match self.ok_or_terminate(message.into_inner(), ctx) {
             Some(recs) => recs,
             None => {
                 tracing::info!("cancelled fetch log task");
                 return;
             }
         };
-        tracing::info!("Pulled Records: {:?}", records.len());
-        let final_record_pulled = if !records.is_empty() {
-            records.get(records.len() - 1)
-        } else {
-            None
-        };
-        match final_record_pulled {
-            Some(record) => {
-                self.pulled_log_offset = Some(record.log_offset);
+        tracing::info!("Pulled Records: {}", output.len());
+        match output.iter().last() {
+            Some((rec, _)) => {
+                self.pulled_log_offset = rec.log_offset;
                 tracing::info!("Pulled Logs Up To Offset: {:?}", self.pulled_log_offset);
-                self.partition(records, ctx).await;
             }
             None => {
                 self.terminate_with_result(
                     Err(CompactionError::InvariantViolation(
-                        "No records pulled by compaction, this is a system invariant violation",
+                        "Logs should be present for compaction",
                     )),
                     ctx,
                 );
+                return;
             }
         }
+        self.partition(output, ctx).await;
+    }
+}
+
+#[async_trait]
+impl Handler<TaskResult<SourceRecordSegmentOutput, SourceRecordSegmentError>>
+    for CompactOrchestrator
+{
+    type Result = ();
+
+    async fn handle(
+        &mut self,
+        message: TaskResult<SourceRecordSegmentOutput, SourceRecordSegmentError>,
+        ctx: &ComponentContext<CompactOrchestrator>,
+    ) {
+        let output = match self.ok_or_terminate(message.into_inner(), ctx) {
+            Some(output) => output,
+            None => return,
+        };
+        tracing::info!("Sourced Records: {}", output.len());
+        self.partition(output, ctx).await;
     }
 }
 
@@ -734,11 +753,11 @@ impl Handler<TaskResult<PartitionOutput, PartitionError>> for CompactOrchestrato
         message: TaskResult<PartitionOutput, PartitionError>,
         ctx: &ComponentContext<CompactOrchestrator>,
     ) {
-        let records = match self.ok_or_terminate(message.into_inner(), ctx) {
+        let output = match self.ok_or_terminate(message.into_inner(), ctx) {
             Some(recs) => recs.records,
             None => todo!(),
         };
-        self.materialize_log(records, ctx).await;
+        self.materialize_log(output, ctx).await;
     }
 }
 
@@ -764,11 +783,11 @@ impl Handler<TaskResult<MaterializeLogOutput, MaterializeLogOperatorError>>
                 && self.num_uncompleted_tasks_by_segment.is_empty()
             {
                 // There is nothing to flush, proceed to register
-                self.register(self.pulled_log_offset.unwrap(), ctx).await;
+                self.register(ctx).await;
             }
         } else {
             self.collection_logical_size_bytes += output.collection_logical_size_delta;
-            self.dispatch_apply_log_to_segment_writer_tasks(output.result, ctx.receiver(), ctx)
+            self.dispatch_apply_log_to_segment_writer_tasks(output.result, ctx)
                 .await;
         }
 
@@ -819,7 +838,7 @@ impl Handler<TaskResult<ApplyLogToSegmentWriterOutput, ApplyLogToSegmentWriterOp
                 None => return,
             };
 
-            self.dispatch_segment_writer_commit(segment_writer, ctx.receiver(), ctx)
+            self.dispatch_segment_writer_commit(segment_writer, ctx)
                 .await;
         }
     }
@@ -841,14 +860,12 @@ impl Handler<TaskResult<CommitSegmentWriterOutput, CommitSegmentWriterOperatorEr
             None => return,
         };
 
-        let flusher = message.flusher;
         // If the flusher recieved is a record segment flusher, get the number of keys for the blockfile and set it on the orchestrator
-        if let ChromaSegmentFlusher::RecordSegment(ref record_segment_flusher) = flusher {
+        if let ChromaSegmentFlusher::RecordSegment(record_segment_flusher) = &message.flusher {
             self.total_records_last_compaction = record_segment_flusher.count();
         }
 
-        self.dispatch_segment_flush(flusher, ctx.receiver(), ctx)
-            .await;
+        self.dispatch_segment_flush(message.flusher, ctx).await;
     }
 }
 
@@ -877,8 +894,7 @@ impl Handler<TaskResult<FlushSegmentWriterOutput, FlushSegmentWriterOperatorErro
         self.num_uncompleted_tasks_by_segment.remove(&segment_id);
 
         if self.num_uncompleted_tasks_by_segment.is_empty() {
-            // Unwrap should be safe here as we are guaranteed to have a value by construction
-            self.register(self.pulled_log_offset.expect("Invariant violation: pulled_log_offset should have been populated at this point."), ctx).await;
+            self.register(ctx).await;
         }
     }
 }
diff --git a/rust/worker/src/execution/orchestration/mod.rs b/rust/worker/src/execution/orchestration/mod.rs
index 91024f40bdb..c013693b23e 100644
--- a/rust/worker/src/execution/orchestration/mod.rs
+++ b/rust/worker/src/execution/orchestration/mod.rs
@@ -1,9 +1,9 @@
-mod compact;
 mod count;
 pub mod spann_knn;
 pub(crate) use compact::*;
 pub(crate) use count::*;
 
+mod compact;
 pub mod get;
 pub mod knn;
 pub mod knn_filter;


Reviews:
user: login: codetheweb
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T22:10:56Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T22:17:23Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T22:18:00Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T22:18:19Z
user: login: codetheweb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T22:20:20Z

Review comments:
diff_hunk: @@ -0,0 +1,80 @@
+use async_trait::async_trait;
+use chroma_error::{ChromaError, ErrorCodes};
+use chroma_segment::blockfile_record::RecordSegmentReader;
+use chroma_system::Operator;
+use chroma_types::{Chunk, LogRecord, OperationRecord};
+use futures::{StreamExt, TryStreamExt};
+use thiserror::Error;
+
+/// The `SourceRecordSegmentOperator` generates logs from the data in the record segment
+///
+/// # Parameters
+/// None
+///
+/// # Inputs
+/// - `record_reader`: The record segment reader, if the collection is initialized
+///
+/// # Outputs
+/// - Chunk of addition logs
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentOperator {}
+
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentInput {
+    pub record_segment_reader: Option<RecordSegmentReader<'static>>,
+}
+
+pub type SourceRecordSegmentOutput = Chunk<LogRecord>;
+
+#[derive(Debug, Error)]
+pub enum SourceRecordSegmentError {
+    #[error("Error reading record segment: {0}")]
+    RecordSegment(#[from] Box<dyn ChromaError>),
+}
+
+impl ChromaError for SourceRecordSegmentError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            SourceRecordSegmentError::RecordSegment(e) => e.code(),
+        }
+    }
+}
+
+#[async_trait]
+impl Operator<SourceRecordSegmentInput, SourceRecordSegmentOutput> for SourceRecordSegmentOperator {
+    type Error = SourceRecordSegmentError;
+    async fn run(
+        &self,
+        input: &SourceRecordSegmentInput,
+    ) -> Result<SourceRecordSegmentOutput, SourceRecordSegmentError> {
+        tracing::trace!("[{}]: {:?}", self.get_name(), input);
+        let logs = match input.record_segment_reader.as_ref() {
+            Some(reader) => {
+                reader
+                    .get_data_stream(..)
+                    .await
+                    .enumerate()
+                    .map(|(offset, res)| {
+                        res.map(|rec| LogRecord {
+                            // Log offset starts with 1
+                            log_offset: offset as i64 + 1,
+                            record: OperationRecord {
+                                id: rec.id.to_string(),
+                                embedding: Some(rec.embedding.to_vec()),
+                                encoding: Some(chroma_types::ScalarEncoding::FLOAT32),
+                                metadata: rec.metadata.map(|meta| {
+                                    meta.into_iter().map(|(k, v)| (k, v.into())).collect()
+                                }),
+                                document: rec.document.map(ToString::to_string),
+                                operation: chroma_types::Operation::Add,
+                            },
+                        })
+                    })
+                    .try_collect::<Vec<_>>()
+                    .await?
+            }
+            None => Default::default(),
+        };
+        Ok(Chunk::new(logs.into()))
+    }
+}
path: rust/worker/src/execution/operators/source_record_segment.rs
user: login: codetheweb
body: might be good to have a basic test
created_at: 2025-04-03T22:06:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -651,20 +642,30 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             self.send(task, ctx).await;
         }
 
-        let fetch_log_task = wrap(
-            Box::new(FetchLogOperator {
-                log_client: self.log.clone(),
-                batch_size: 100,
-                // Here we do not need to be inclusive since the compaction job
-                // offset is the one after the last compaction offset
-                start_log_offset_id: self.compaction_job.offset as u32,
-                maximum_fetch_count: Some(self.max_compaction_size as u32),
-                collection_uuid: self.compaction_job.collection_id,
-            }),
-            (),
-            ctx.receiver(),
-        );
-        self.send(fetch_log_task, ctx).await;
+        let log_task = match self.compaction_job.rebuild {
path: rust/worker/src/execution/orchestration/compact.rs
user: login: codetheweb
body: nit: if we're rebuilding we could skip the prefetch tasks
created_at: 2025-04-03T22:10:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -651,20 +642,30 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             self.send(task, ctx).await;
         }
 
-        let fetch_log_task = wrap(
-            Box::new(FetchLogOperator {
-                log_client: self.log.clone(),
-                batch_size: 100,
-                // Here we do not need to be inclusive since the compaction job
-                // offset is the one after the last compaction offset
-                start_log_offset_id: self.compaction_job.offset as u32,
-                maximum_fetch_count: Some(self.max_compaction_size as u32),
-                collection_uuid: self.compaction_job.collection_id,
-            }),
-            (),
-            ctx.receiver(),
-        );
-        self.send(fetch_log_task, ctx).await;
+        let log_task = match self.compaction_job.rebuild {
path: rust/worker/src/execution/orchestration/compact.rs
user: login: Sicheng-Pan
body: We still need to prefetch the record segment right?
created_at: 2025-04-03T22:17:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,80 @@
+use async_trait::async_trait;
+use chroma_error::{ChromaError, ErrorCodes};
+use chroma_segment::blockfile_record::RecordSegmentReader;
+use chroma_system::Operator;
+use chroma_types::{Chunk, LogRecord, OperationRecord};
+use futures::{StreamExt, TryStreamExt};
+use thiserror::Error;
+
+/// The `SourceRecordSegmentOperator` generates logs from the data in the record segment
+///
+/// # Parameters
+/// None
+///
+/// # Inputs
+/// - `record_reader`: The record segment reader, if the collection is initialized
+///
+/// # Outputs
+/// - Chunk of addition logs
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentOperator {}
+
+#[derive(Clone, Debug)]
+pub struct SourceRecordSegmentInput {
+    pub record_segment_reader: Option<RecordSegmentReader<'static>>,
+}
+
+pub type SourceRecordSegmentOutput = Chunk<LogRecord>;
+
+#[derive(Debug, Error)]
+pub enum SourceRecordSegmentError {
+    #[error("Error reading record segment: {0}")]
+    RecordSegment(#[from] Box<dyn ChromaError>),
+}
+
+impl ChromaError for SourceRecordSegmentError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            SourceRecordSegmentError::RecordSegment(e) => e.code(),
+        }
+    }
+}
+
+#[async_trait]
+impl Operator<SourceRecordSegmentInput, SourceRecordSegmentOutput> for SourceRecordSegmentOperator {
+    type Error = SourceRecordSegmentError;
+    async fn run(
+        &self,
+        input: &SourceRecordSegmentInput,
+    ) -> Result<SourceRecordSegmentOutput, SourceRecordSegmentError> {
+        tracing::trace!("[{}]: {:?}", self.get_name(), input);
+        let logs = match input.record_segment_reader.as_ref() {
+            Some(reader) => {
+                reader
+                    .get_data_stream(..)
+                    .await
+                    .enumerate()
+                    .map(|(offset, res)| {
+                        res.map(|rec| LogRecord {
+                            // Log offset starts with 1
+                            log_offset: offset as i64 + 1,
+                            record: OperationRecord {
+                                id: rec.id.to_string(),
+                                embedding: Some(rec.embedding.to_vec()),
+                                encoding: Some(chroma_types::ScalarEncoding::FLOAT32),
+                                metadata: rec.metadata.map(|meta| {
+                                    meta.into_iter().map(|(k, v)| (k, v.into())).collect()
+                                }),
+                                document: rec.document.map(ToString::to_string),
+                                operation: chroma_types::Operation::Add,
+                            },
+                        })
+                    })
+                    .try_collect::<Vec<_>>()
+                    .await?
+            }
+            None => Default::default(),
+        };
+        Ok(Chunk::new(logs.into()))
+    }
+}
path: rust/worker/src/execution/operators/source_record_segment.rs
user: login: Sicheng-Pan
body: Will add a unit test for this
created_at: 2025-04-03T22:18:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -651,20 +642,30 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             self.send(task, ctx).await;
         }
 
-        let fetch_log_task = wrap(
-            Box::new(FetchLogOperator {
-                log_client: self.log.clone(),
-                batch_size: 100,
-                // Here we do not need to be inclusive since the compaction job
-                // offset is the one after the last compaction offset
-                start_log_offset_id: self.compaction_job.offset as u32,
-                maximum_fetch_count: Some(self.max_compaction_size as u32),
-                collection_uuid: self.compaction_job.collection_id,
-            }),
-            (),
-            ctx.receiver(),
-        );
-        self.send(fetch_log_task, ctx).await;
+        let log_task = match self.compaction_job.rebuild {
path: rust/worker/src/execution/orchestration/compact.rs
user: login: Sicheng-Pan
body: The metadata segment does not be prefetched during rebuild
created_at: 2025-04-03T22:18:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -651,20 +642,30 @@ impl Handler<TaskResult<GetCollectionAndSegmentsOutput, GetCollectionAndSegments
             self.send(task, ctx).await;
         }
 
-        let fetch_log_task = wrap(
-            Box::new(FetchLogOperator {
-                log_client: self.log.clone(),
-                batch_size: 100,
-                // Here we do not need to be inclusive since the compaction job
-                // offset is the one after the last compaction offset
-                start_log_offset_id: self.compaction_job.offset as u32,
-                maximum_fetch_count: Some(self.max_compaction_size as u32),
-                collection_uuid: self.compaction_job.collection_id,
-            }),
-            (),
-            ctx.receiver(),
-        );
-        self.send(fetch_log_task, ctx).await;
+        let log_task = match self.compaction_job.rebuild {
path: rust/worker/src/execution/orchestration/compact.rs
user: login: codetheweb
body: ah yep
created_at: 2025-04-03T22:20:20Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-02T00:35:05Z
body: # Reviewer Checklist
  Please leverage this checklist to ensure your code review is thorough before approving
  ## Testing, Bugs, Errors, Logs, Documentation
  - [ ] Can you think of any use case in which the code does not behave as intended? Have they been tested?
  - [ ] Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  - [ ] If appropriate, are there adequate property based tests?
  - [ ] If appropriate, are there adequate unit tests?
  - [ ] Should any logging, debugging, tracing information be added or removed?
  - [ ] Are error messages user-friendly?
  - [ ] Have all documentation changes needed been made?
  - [ ] Have all non-obvious changes been commented?
  ## System Compatibility
  - [ ] Are there any potential impacts on other parts of the system or backward compatibility?
  - [ ] Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?
  ## Quality
  - [ ] Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)
author_association: NONE
user: login: Sicheng-Pan
created_at: 2025-04-02T00:35:10Z
body: * **#4132** <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4132?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a>
* **#4127** <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4127?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a> 👈 <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4127?utm_source=stack-comment-view-in-graphite" target="_blank">(View in Graphite)</a>
* **#4125** <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4125?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a>
* `main`


<h2></h2>

This stack of pull requests is managed by <a href="https://graphite.dev?utm-source=stack-comment"><b>Graphite</b></a>. Learn more about <a href="https://stacking.dev/?utm_source=stack-comment">stacking</a>.
<!-- Current dependencies on/for this PR: -->
author_association: CONTRIBUTOR
user: login: Sicheng-Pan
created_at: 2025-04-04T22:42:36Z
body: Note: the failing test will be fixed in the top PR
author_association: CONTRIBUTOR
user: login: Sicheng-Pan
created_at: 2025-04-04T23:42:26Z
body: ### Merge activity

* **Apr 4, 7:42 PM EDT**: A user started a stack merge that includes this pull request via [Graphite](https://app.graphite.dev/github/pr/chroma-core/chroma/4127).
* **Apr 4, 7:44 PM EDT**: [Graphite](https://app.graphite.dev/github/pr/chroma-core/chroma/4127) rebased this pull request as part of a merge.
* **Apr 4, 7:45 PM EDT**: A user merged this pull request with [Graphite](https://app.graphite.dev/github/pr/chroma-core/chroma/4127).

author_association: CONTRIBUTOR