Make data_record owned

codetheweb · codetheweb · commit e198444824b7 · 2024-12-02T17:16:31.000-08:00
diff --git a/rust/types/src/data_record.rs b/rust/types/src/data_record.rs
@@ -10,6 +10,29 @@ pub struct DataRecord<'a> {
     pub document: Option<&'a str>,
 }
 
+#[derive(Debug, Clone)]
+pub struct OwnedDataRecord {
+    pub id: String,
+    pub embedding: Vec<f32>,
+    pub metadata: Option<Metadata>,
+    pub document: Option<String>,
+}
+
+impl<'a> From<&DataRecord<'a>> for OwnedDataRecord {
+    fn from(data_record: &DataRecord<'a>) -> Self {
+        let id = data_record.id.to_string();
+        let embedding = data_record.embedding.to_vec();
+        let metadata = data_record.metadata.clone();
+        let document = data_record.document.map(|doc| doc.to_string());
+        OwnedDataRecord {
+            id,
+            embedding,
+            metadata,
+            document,
+        }
+    }
+}
+
 impl DataRecord<'_> {
     pub fn get_size(&self) -> usize {
         let id_size = self.id.len();
@@ -28,4 +51,8 @@ impl DataRecord<'_> {
         };
         id_size + embedding_size + metadata_size + document_size
     }
+
+    pub fn to_owned(&self) -> OwnedDataRecord {
+        self.into()
+    }
 }
diff --git a/rust/worker/src/execution/operators/filter.rs b/rust/worker/src/execution/operators/filter.rs
@@ -106,7 +106,7 @@ pub(crate) struct MetadataLogReader<'me> {
 }
 
 impl<'me> MetadataLogReader<'me> {
-    pub(crate) fn new(logs: &'me Chunk<MaterializedLogRecord<'me>>) -> Self {
+    pub(crate) fn new(logs: &'me Chunk<MaterializedLogRecord>) -> Self {
         let mut compact_metadata: HashMap<_, BTreeMap<&MetadataValue, RoaringBitmap>> =
             HashMap::new();
         let mut document = HashMap::new();
diff --git a/rust/worker/src/execution/operators/hnsw_knn.rs b/rust/worker/src/execution/operators/hnsw_knn.rs
@@ -65,7 +65,7 @@ impl ChromaError for HnswKnnOperatorError {
 impl HnswKnnOperator {
     async fn get_disallowed_ids<'referred_data>(
         &self,
-        logs: Chunk<MaterializedLogRecord<'_>>,
+        logs: Chunk<MaterializedLogRecord>,
         record_segment_reader: &RecordSegmentReader<'_>,
     ) -> Result<Vec<u32>, Box<dyn ChromaError>> {
         let mut disallowed_ids = Vec::new();
diff --git a/rust/worker/src/segment/distributed_hnsw_segment.rs b/rust/worker/src/segment/distributed_hnsw_segment.rs
@@ -238,10 +238,10 @@ impl DistributedHNSWSegmentWriter {
     }
 }
 
-impl<'a> SegmentWriter<'a> for DistributedHNSWSegmentWriter {
+impl SegmentWriter for DistributedHNSWSegmentWriter {
     async fn apply_materialized_log_chunk(
         &self,
-        records: chroma_types::Chunk<super::MaterializedLogRecord<'a>>,
+        records: chroma_types::Chunk<super::MaterializedLogRecord>,
     ) -> Result<(), ApplyMaterializedLogError> {
         for (record, _) in records.iter() {
             match record.final_operation {
diff --git a/rust/worker/src/segment/metadata_segment.rs b/rust/worker/src/segment/metadata_segment.rs
@@ -530,15 +530,19 @@ impl<'me> MetadataSegmentWriter<'me> {
     }
 }
 
-impl<'log_records> SegmentWriter<'log_records> for MetadataSegmentWriter<'_> {
+impl SegmentWriter for MetadataSegmentWriter<'_> {
     async fn apply_materialized_log_chunk(
         &self,
-        records: Chunk<MaterializedLogRecord<'log_records>>,
+        records: Chunk<MaterializedLogRecord>,
     ) -> Result<(), ApplyMaterializedLogError> {
         let mut count = 0u64;
         let full_text_writer_batch = records.iter().filter_map(|record| {
             let offset_id = record.0.offset_id;
-            let old_document = record.0.data_record.as_ref().and_then(|r| r.document);
+            let old_document = record
+                .0
+                .data_record
+                .as_ref()
+                .and_then(|r| r.document.as_ref().map(|d| d.as_str()));
             let new_document = &record.0.final_document;
 
             if matches!(
diff --git a/rust/worker/src/segment/record_segment.rs b/rust/worker/src/segment/record_segment.rs
@@ -61,9 +61,9 @@ pub enum RecordSegmentWriterCreationError {
 }
 
 impl RecordSegmentWriter {
-    async fn construct_and_set_data_record<'a>(
+    async fn construct_and_set_data_record(
         &self,
-        mat_record: &MaterializedLogRecord<'a>,
+        mat_record: &MaterializedLogRecord,
         user_id: &str,
         offset_id: u32,
     ) -> Result<(), ApplyMaterializedLogError> {
@@ -337,10 +337,10 @@ impl ChromaError for ApplyMaterializedLogError {
     }
 }
 
-impl<'a> SegmentWriter<'a> for RecordSegmentWriter {
+impl SegmentWriter for RecordSegmentWriter {
     async fn apply_materialized_log_chunk(
         &self,
-        records: Chunk<MaterializedLogRecord<'a>>,
+        records: Chunk<MaterializedLogRecord>,
     ) -> Result<(), ApplyMaterializedLogError> {
         // The max new offset id introduced by materialized logs is initialized as zero
         // Since offset id should start from 1, we use this to indicate no new offset id
@@ -423,7 +423,7 @@ impl<'a> SegmentWriter<'a> for RecordSegmentWriter {
                     match self
                         .construct_and_set_data_record(
                             log_record,
-                            log_record.data_record.as_ref().unwrap().id,
+                            &log_record.data_record.as_ref().unwrap().id,
                             log_record.offset_id,
                         )
                         .await
@@ -440,7 +440,7 @@ impl<'a> SegmentWriter<'a> for RecordSegmentWriter {
                         .user_id_to_id
                         .as_ref()
                         .unwrap()
-                        .delete::<&str, u32>("", log_record.data_record.as_ref().unwrap().id)
+                        .delete::<&str, u32>("", &log_record.data_record.as_ref().unwrap().id)
                         .await
                     {
                         Ok(()) => (),
diff --git a/rust/worker/src/segment/types.rs b/rust/worker/src/segment/types.rs
@@ -3,7 +3,7 @@ use chroma_error::{ChromaError, ErrorCodes};
 use chroma_types::{
     Chunk, DataRecord, DeletedMetadata, LogRecord, MaterializedLogOperation, Metadata,
     MetadataDelta, MetadataValue, MetadataValueConversionError, Operation, OperationRecord,
-    UpdateMetadata, UpdateMetadataValue,
+    OwnedDataRecord, UpdateMetadata, UpdateMetadataValue,
 };
 use std::collections::{HashMap, HashSet};
 use std::sync::atomic::AtomicU32;
@@ -113,10 +113,10 @@ impl ChromaError for LogMaterializerError {
 }
 
 #[derive(Debug, Clone)]
-pub struct MaterializedLogRecord<'referred_data> {
+pub struct MaterializedLogRecord {
     // This is the data record read from the record segment for this id.
     // None if the record exists only in the log.
-    pub(crate) data_record: Option<DataRecord<'referred_data>>,
+    pub(crate) data_record: Option<OwnedDataRecord>,
     // If present in the record segment then it is the offset id
     // in the record segment at which the record was found.
     // If not present in the segment then it is the offset id
@@ -157,7 +157,7 @@ pub struct MaterializedLogRecord<'referred_data> {
     pub(crate) final_embedding: Option<Vec<f32>>,
 }
 
-impl<'referred_data> MaterializedLogRecord<'referred_data> {
+impl MaterializedLogRecord {
     // Performs a deep copy of the document so only use it if really
     // needed. If you only need a reference then use merged_document_ref
     // defined below.
@@ -170,7 +170,7 @@ impl<'referred_data> MaterializedLogRecord<'referred_data> {
         return match self.final_document.clone() {
             Some(doc) => Some(doc),
             None => match self.data_record.as_ref() {
-                Some(data_record) => data_record.document.map(|doc| doc.to_string()),
+                Some(data_record) => data_record.document.clone(),
                 None => None,
             },
         };
@@ -185,10 +185,7 @@ impl<'referred_data> MaterializedLogRecord<'referred_data> {
         return match &self.final_document {
             Some(doc) => Some(doc),
             None => match self.data_record.as_ref() {
-                Some(data_record) => match data_record.document {
-                    Some(doc) => Some(doc),
-                    None => None,
-                },
+                Some(data_record) => data_record.document.as_deref(),
                 None => None,
             },
         };
@@ -211,7 +208,7 @@ impl<'referred_data> MaterializedLogRecord<'referred_data> {
         match &self.user_id {
             Some(id) => id.as_str(),
             None => match &self.data_record {
-                Some(data_record) => data_record.id,
+                Some(data_record) => &data_record.id,
                 None => panic!("Expected at least one user id to be set"),
             },
         }
@@ -247,7 +244,7 @@ impl<'referred_data> MaterializedLogRecord<'referred_data> {
         final_metadata
     }
 
-    pub(crate) fn metadata_delta(&'referred_data self) -> MetadataDelta<'referred_data> {
+    pub(crate) fn metadata_delta(&self) -> MetadataDelta<'_> {
         let mut metadata_delta = MetadataDelta::new();
         let mut base_metadata: HashMap<&str, &MetadataValue> = HashMap::new();
         if let Some(data_record) = &self.data_record {
@@ -327,21 +324,19 @@ impl<'referred_data> MaterializedLogRecord<'referred_data> {
         return match &self.final_embedding {
             Some(embed) => embed,
             None => match self.data_record.as_ref() {
-                Some(data_record) => data_record.embedding,
+                Some(data_record) => &data_record.embedding,
                 None => panic!("Expected at least one source of embedding"),
             },
         };
     }
 }
 
-impl<'referred_data> From<(DataRecord<'referred_data>, u32)>
-    for MaterializedLogRecord<'referred_data>
-{
+impl<'referred_data> From<(DataRecord<'referred_data>, u32)> for MaterializedLogRecord {
     fn from(data_record_info: (DataRecord<'referred_data>, u32)) -> Self {
         let data_record = data_record_info.0;
         let offset_id = data_record_info.1;
         Self {
-            data_record: Some(data_record),
+            data_record: Some(data_record.to_owned()),
             offset_id,
             user_id: None,
             final_operation: MaterializedLogOperation::Initial,
@@ -357,7 +352,7 @@ impl<'referred_data> From<(DataRecord<'referred_data>, u32)>
 // in the log (OperationRecord), offset id in storage where it will be stored (u32)
 // and user id (str).
 impl<'referred_data> TryFrom<(&'referred_data OperationRecord, u32, &'referred_data str)>
-    for MaterializedLogRecord<'referred_data>
+    for MaterializedLogRecord
 {
     type Error = LogMaterializerError;
 
@@ -414,7 +409,7 @@ pub async fn materialize_logs<'me>(
     // for materializing. Writers pass this value to the materializer
     // because they need to share this across all log partitions.
     next_offset_id: Option<Arc<AtomicU32>>,
-) -> Result<Chunk<MaterializedLogRecord<'me>>, LogMaterializerError> {
+) -> Result<Chunk<MaterializedLogRecord>, LogMaterializerError> {
     // Trace the total_len since len() iterates over the entire chunk
     // and we don't want to do that just to trace the length.
     tracing::info!("Total length of logs in materializer: {}", logs.total_len());
@@ -748,10 +743,10 @@ pub async fn materialize_logs<'me>(
 
 // This needs to be public for testing
 #[allow(async_fn_in_trait)]
-pub trait SegmentWriter<'a> {
+pub trait SegmentWriter {
     async fn apply_materialized_log_chunk(
         &self,
-        records: Chunk<MaterializedLogRecord<'a>>,
+        records: Chunk<MaterializedLogRecord>,
     ) -> Result<(), ApplyMaterializedLogError>;
     async fn commit(self) -> Result<impl SegmentFlusher, Box<dyn ChromaError>>;
 }
@@ -1880,7 +1875,10 @@ mod tests {
                 assert_eq!(hello_found, 1);
                 assert_eq!(hello_again_found, 1);
                 assert!(log.data_record.is_some());
-                assert_eq!(log.data_record.as_ref().unwrap().document, Some("doc1"));
+                assert_eq!(
+                    log.data_record.as_ref().unwrap().document,
+                    Some("doc1".to_string())
+                );
                 assert_eq!(
                     log.data_record.as_ref().unwrap().embedding,
                     vec![1.0, 2.0, 3.0].as_slice()

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ pub(crate) struct MetadataLogReader<'me> {`
`106`	`106`	`}`
`107`	`107`
`108`	`108`	`impl<'me> MetadataLogReader<'me> {`
`109`		`- pub(crate) fn new(logs: &'me Chunk<MaterializedLogRecord<'me>>) -> Self {`
	`109`	`+ pub(crate) fn new(logs: &'me Chunk<MaterializedLogRecord>) -> Self {`
`110`	`110`	`let mut compact_metadata: HashMap<_, BTreeMap<&MetadataValue, RoaringBitmap>> =`
`111`	`111`	`HashMap::new();`
`112`	`112`	`let mut document = HashMap::new();`
Original file line number	Diff line number	Diff line change
`@@ -238,10 +238,10 @@ impl DistributedHNSWSegmentWriter {`
`238`	`238`	`}`
`239`	`239`	`}`
`240`	`240`
`241`		`-impl<'a> SegmentWriter<'a> for DistributedHNSWSegmentWriter {`
	`241`	`+impl SegmentWriter for DistributedHNSWSegmentWriter {`
`242`	`242`	`async fn apply_materialized_log_chunk(`
`243`	`243`	`&self,`
`244`		`- records: chroma_types::Chunk<super::MaterializedLogRecord<'a>>,`
	`244`	`+ records: chroma_types::Chunk<super::MaterializedLogRecord>,`
`245`	`245`	`) -> Result<(), ApplyMaterializedLogError> {`
`246`	`246`	`for (record, _) in records.iter() {`
`247`	`247`	`match record.final_operation {`