[CLN]: make materialization function rather than struct (#3165)

## Description of changes Simplifies usage of log materialization logic & lifetimes in preparation for next PR in this stack. There are no actual logic changes. ## 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 *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)?* n/a
chroma-core · Dec 2, 2024 · ccff355 · ccff355
1 parent 337fe73
commit ccff355
Show file tree

Hide file tree

Showing 14 changed files with 160 additions and 243 deletions.
diff --git a/rust/worker/src/execution/operators/brute_force_knn.rs b/rust/worker/src/execution/operators/brute_force_knn.rs
@@ -1,8 +1,7 @@
 use crate::execution::operator::Operator;
 use crate::execution::operators::normalize_vectors::normalize;
 use crate::segment::record_segment::RecordSegmentReader;
-use crate::segment::LogMaterializer;
-use crate::segment::LogMaterializerError;
+use crate::segment::{materialize_logs, LogMaterializerError};
 use async_trait::async_trait;
 use chroma_blockstore::provider::BlockfileProvider;
 use chroma_distance::DistanceFunction;
@@ -134,9 +133,7 @@ impl Operator<BruteForceKnnOperatorInput, BruteForceKnnOperatorOutput> for Brute
                 }
             }
         };
-        let log_materializer = LogMaterializer::new(record_segment_reader, input.log.clone(), None);
-        let logs = match log_materializer
-            .materialize()
+        let logs = match materialize_logs(&record_segment_reader, &input.log, None)
             .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
             .await
         {

diff --git a/rust/worker/src/execution/operators/count_records.rs b/rust/worker/src/execution/operators/count_records.rs
@@ -199,9 +199,9 @@ impl Operator<CountRecordsInput, CountRecordsOutput> for CountRecordsOperator {
 
 #[cfg(test)]
 mod tests {
+    use crate::segment::materialize_logs;
     use crate::segment::record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError};
     use crate::segment::types::SegmentFlusher;
-    use crate::segment::LogMaterializer;
     use crate::{
         execution::{
             operator::Operator,
@@ -293,9 +293,7 @@ mod tests {
                         }
                     }
                 };
-            let materializer = LogMaterializer::new(record_segment_reader, data, None);
-            let mat_records = materializer
-                .materialize()
+            let mat_records = materialize_logs(&record_segment_reader, &data, None)
                 .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
                 .await
                 .expect("Log materialization failed");

diff --git a/rust/worker/src/execution/operators/filter.rs b/rust/worker/src/execution/operators/filter.rs
@@ -19,9 +19,10 @@ use tracing::{trace, Instrument, Span};
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         metadata_segment::{MetadataSegmentError, MetadataSegmentReader},
         record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
-        LogMaterializer, LogMaterializerError, MaterializedLogRecord,
+        LogMaterializerError, MaterializedLogRecord,
     },
 };
 
@@ -416,10 +417,8 @@ impl Operator<FilterInput, FilterOutput> for FilterOperator {
             }
             Err(e) => Err(*e),
         }?;
-        let materializer =
-            LogMaterializer::new(record_segment_reader.clone(), input.logs.clone(), None);
-        let materialized_logs = materializer
-            .materialize()
+        let cloned_record_segment_reader = record_segment_reader.clone();
+        let materialized_logs = materialize_logs(&cloned_record_segment_reader, &input.logs, None)
             .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
             .await?;
         let metadata_log_reader = MetadataLogReader::new(&materialized_logs);

diff --git a/rust/worker/src/execution/operators/get_vectors_operator.rs b/rust/worker/src/execution/operators/get_vectors_operator.rs
@@ -1,8 +1,9 @@
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         record_segment::{self, RecordSegmentReader},
-        LogMaterializer, LogMaterializerError,
+        LogMaterializerError,
     },
 };
 use async_trait::async_trait;
@@ -120,17 +121,13 @@ impl Operator<GetVectorsOperatorInput, GetVectorsOperatorOutput> for GetVectorsO
             },
         };
         // Step 1: Materialize the logs.
-        let materializer = LogMaterializer::new(
-            record_segment_reader.clone(),
-            input.log_records.clone(),
-            None,
-        );
-        let mat_records = match materializer.materialize().await {
-            Ok(records) => records,
-            Err(e) => {
-                return Err(GetVectorsOperatorError::LogMaterialization(e));
-            }
-        };
+        let mat_records =
+            match materialize_logs(&record_segment_reader, &input.log_records, None).await {
+                Ok(records) => records,
+                Err(e) => {
+                    return Err(GetVectorsOperatorError::LogMaterialization(e));
+                }
+            };
 
         // Search the log records for the user ids
         let mut remaining_search_user_ids: HashSet<String> =

diff --git a/rust/worker/src/execution/operators/hnsw_knn.rs b/rust/worker/src/execution/operators/hnsw_knn.rs
@@ -1,5 +1,5 @@
 use crate::segment::record_segment::RecordSegmentReaderCreationError;
-use crate::segment::{LogMaterializer, LogMaterializerError, MaterializedLogRecord};
+use crate::segment::{materialize_logs, LogMaterializerError, MaterializedLogRecord};
 use crate::{
     execution::operator::Operator,
     segment::{
@@ -147,13 +147,8 @@ impl Operator<HnswKnnOperatorInput, HnswKnnOperatorOutput> for HnswKnnOperator {
                 }
             },
         };
-        let log_materializer = LogMaterializer::new(
-            Some(record_segment_reader.clone()),
-            input.logs.clone(),
-            None,
-        );
-        let logs = match log_materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader.clone());
+        let logs = match materialize_logs(&some_reader, &input.logs, None)
             .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
             .await
         {

diff --git a/rust/worker/src/execution/operators/knn_log.rs b/rust/worker/src/execution/operators/knn_log.rs
@@ -10,8 +10,9 @@ use tonic::async_trait;
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
-        LogMaterializer, LogMaterializerError,
+        LogMaterializerError,
     },
 };
 
@@ -72,8 +73,7 @@ impl Operator<KnnLogInput, KnnLogOutput> for KnnOperator {
             Err(e) => Err(*e),
         }?;
 
-        let materializer = LogMaterializer::new(record_segment_reader, input.logs.clone(), None);
-        let logs = materializer.materialize().await?;
+        let logs = materialize_logs(&record_segment_reader, &input.logs, None).await?;
 
         let target_vector;
         let target_embedding = if let DistanceFunction::Cosine = input.distance_function {

diff --git a/rust/worker/src/execution/operators/limit.rs b/rust/worker/src/execution/operators/limit.rs
@@ -11,8 +11,9 @@ use tracing::{trace, Instrument, Span};
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
-        LogMaterializer, LogMaterializerError,
+        LogMaterializerError,
     },
 };
 
@@ -213,10 +214,7 @@ impl Operator<LimitInput, LimitOutput> for LimitOperator {
         let mut materialized_log_offset_ids = match &input.log_offset_ids {
             SignedRoaringBitmap::Include(rbm) => rbm.clone(),
             SignedRoaringBitmap::Exclude(rbm) => {
-                let materializer =
-                    LogMaterializer::new(record_segment_reader.clone(), input.logs.clone(), None);
-                let materialized_logs = materializer
-                    .materialize()
+                let materialized_logs = materialize_logs(&record_segment_reader, &input.logs, None)
                     .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
                     .await?;
 

diff --git a/rust/worker/src/execution/operators/prefetch_record.rs b/rust/worker/src/execution/operators/prefetch_record.rs
@@ -10,8 +10,9 @@ use tracing::{trace, Instrument, Span};
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
-        LogMaterializer, LogMaterializerError,
+        LogMaterializerError,
     },
 };
 
@@ -84,13 +85,8 @@ impl Operator<PrefetchRecordInput, PrefetchRecordOutput> for PrefetchRecordOpera
             Err(e) => return Err((*e).into()),
         };
 
-        let materializer = LogMaterializer::new(
-            Some(record_segment_reader.clone()),
-            input.logs.clone(),
-            None,
-        );
-        let materialized_logs = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader.clone());
+        let materialized_logs = materialize_logs(&some_reader, &input.logs, None)
             .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
             .await?;
 

diff --git a/rust/worker/src/execution/operators/projection.rs b/rust/worker/src/execution/operators/projection.rs
@@ -10,8 +10,9 @@ use tracing::{error, trace, Instrument, Span};
 use crate::{
     execution::operator::Operator,
     segment::{
+        materialize_logs,
         record_segment::{RecordSegmentReader, RecordSegmentReaderCreationError},
-        LogMaterializer, LogMaterializerError,
+        LogMaterializerError,
     },
 };
 
@@ -104,10 +105,8 @@ impl Operator<ProjectionInput, ProjectionOutput> for ProjectionOperator {
             }
             Err(e) => Err(*e),
         }?;
-        let materializer =
-            LogMaterializer::new(record_segment_reader.clone(), input.logs.clone(), None);
-        let materialized_logs = materializer
-            .materialize()
+
+        let materialized_logs = materialize_logs(&record_segment_reader, &input.logs, None)
             .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
             .await?;
 

diff --git a/rust/worker/src/execution/operators/write_segments.rs b/rust/worker/src/execution/operators/write_segments.rs
@@ -1,9 +1,9 @@
+use crate::segment::materialize_logs;
 use crate::segment::metadata_segment::MetadataSegmentError;
 use crate::segment::metadata_segment::MetadataSegmentWriter;
 use crate::segment::record_segment::ApplyMaterializedLogError;
 use crate::segment::record_segment::RecordSegmentReader;
 use crate::segment::record_segment::RecordSegmentReaderCreationError;
-use crate::segment::LogMaterializer;
 use crate::segment::LogMaterializerError;
 use crate::segment::SegmentWriter;
 use crate::{
@@ -160,16 +160,14 @@ impl Operator<WriteSegmentsInput, WriteSegmentsOutput> for WriteSegmentsOperator
                 };
             }
         };
-        let materializer = LogMaterializer::new(
-            record_segment_reader,
-            input.chunk.clone(),
-            Some(input.next_offset_id.clone()),
-        );
         // Materialize the logs.
-        let res = match materializer
-            .materialize()
-            .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
-            .await
+        let res = match materialize_logs(
+            &record_segment_reader,
+            &input.chunk,
+            Some(input.next_offset_id.clone()),
+        )
+        .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
+        .await
         {
             Ok(records) => records,
             Err(e) => {

diff --git a/rust/worker/src/segment/metadata_segment.rs b/rust/worker/src/segment/metadata_segment.rs
@@ -1075,11 +1075,12 @@ mod test {
     #![allow(deprecated)]
 
     use crate::segment::{
+        materialize_logs,
         metadata_segment::{MetadataSegmentReader, MetadataSegmentWriter},
         record_segment::{
             RecordSegmentReader, RecordSegmentReaderCreationError, RecordSegmentWriter,
         },
-        LogMaterializer, SegmentFlusher, SegmentWriter,
+        SegmentFlusher, SegmentWriter,
     };
     use chroma_blockstore::{
         arrow::{config::TEST_MAX_BLOCK_SIZE_BYTES, provider::ArrowBlockfileProvider},
@@ -1193,9 +1194,7 @@ mod test {
                         }
                     }
                 };
-            let materializer = LogMaterializer::new(record_segment_reader, data, None);
-            let mat_records = materializer
-                .materialize()
+            let mat_records = materialize_logs(&record_segment_reader, &data, None)
                 .await
                 .expect("Log materialization failed");
             metadata_writer
@@ -1265,9 +1264,8 @@ mod test {
             MetadataSegmentWriter::from_segment(&metadata_segment, &blockfile_provider)
                 .await
                 .expect("Error creating segment writer");
-        let materializer = LogMaterializer::new(Some(record_segment_reader), data, None);
-        let mat_records = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader);
+        let mat_records = materialize_logs(&some_reader, &data, None)
             .await
             .expect("Log materialization failed");
         metadata_writer
@@ -1347,9 +1345,8 @@ mod test {
             MetadataSegmentWriter::from_segment(&metadata_segment, &blockfile_provider)
                 .await
                 .expect("Error creating segment writer");
-        let materializer = LogMaterializer::new(Some(record_segment_reader), data, None);
-        let mat_records = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader);
+        let mat_records = materialize_logs(&some_reader, &data, None)
             .await
             .expect("Log materialization failed");
         metadata_writer
@@ -1487,9 +1484,7 @@ mod test {
                         }
                     }
                 };
-            let materializer = LogMaterializer::new(record_segment_reader, data, None);
-            let mat_records = materializer
-                .materialize()
+            let mat_records = materialize_logs(&record_segment_reader, &data, None)
                 .await
                 .expect("Log materialization failed");
             metadata_writer
@@ -1566,9 +1561,8 @@ mod test {
             MetadataSegmentWriter::from_segment(&metadata_segment, &blockfile_provider)
                 .await
                 .expect("Error creating segment writer");
-        let materializer = LogMaterializer::new(Some(record_segment_reader), data, None);
-        let mat_records = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader);
+        let mat_records = materialize_logs(&some_reader, &data, None)
             .await
             .expect("Log materialization failed");
         metadata_writer
@@ -1740,9 +1734,7 @@ mod test {
                         }
                     }
                 };
-            let materializer = LogMaterializer::new(record_segment_reader, data, None);
-            let mat_records = materializer
-                .materialize()
+            let mat_records = materialize_logs(&record_segment_reader, &data, None)
                 .await
                 .expect("Log materialization failed");
             metadata_writer
@@ -1801,9 +1793,8 @@ mod test {
             MetadataSegmentWriter::from_segment(&metadata_segment, &blockfile_provider)
                 .await
                 .expect("Error creating segment writer");
-        let materializer = LogMaterializer::new(Some(record_segment_reader), data, None);
-        let mat_records = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader);
+        let mat_records = materialize_logs(&some_reader, &data, None)
             .await
             .expect("Log materialization failed");
         metadata_writer
@@ -1962,9 +1953,7 @@ mod test {
                         }
                     }
                 };
-            let materializer = LogMaterializer::new(record_segment_reader, data, None);
-            let mat_records = materializer
-                .materialize()
+            let mat_records = materialize_logs(&record_segment_reader, &data, None)
                 .await
                 .expect("Log materialization failed");
             metadata_writer
@@ -2021,9 +2010,8 @@ mod test {
             MetadataSegmentWriter::from_segment(&metadata_segment, &blockfile_provider)
                 .await
                 .expect("Error creating segment writer");
-        let materializer = LogMaterializer::new(Some(record_segment_reader), data, None);
-        let mat_records = materializer
-            .materialize()
+        let some_reader = Some(record_segment_reader);
+        let mat_records = materialize_logs(&some_reader, &data, None)
             .await
             .expect("Log materialization failed");
         metadata_writer