[ENH]: Make all functions incremental

tanujnay112 · tanujnay112 · commit 1780f2e61076 · 2025-11-21T02:03:34.000-08:00
diff --git a/rust/worker/src/execution/functions/statistics.rs b/rust/worker/src/execution/functions/statistics.rs
@@ -3,7 +3,8 @@
 //! The core idea is the following: For each key-value pair associated with a record, aggregate so
 //! (key, value) -> count.  This gives a count of how frequently each key appears.
 //!
-//! For now it's not incremental.
+//! The statistics executor is incremental - it loads existing counts from the output_reader
+//! and updates them with new records.
 
 use std::collections::{HashMap, HashSet};
 use std::hash::{Hash, Hasher};
@@ -13,8 +14,7 @@ use chroma_error::ChromaError;
 use chroma_segment::blockfile_record::RecordSegmentReader;
 use chroma_segment::types::HydratedMaterializedLogRecord;
 use chroma_types::{
-    Chunk, LogRecord, MaterializedLogOperation, MetadataValue, Operation, OperationRecord,
-    UpdateMetadataValue,
+    Chunk, LogRecord, MetadataValue, Operation, OperationRecord, UpdateMetadataValue,
 };
 use futures::StreamExt;
 
@@ -27,8 +27,10 @@ pub trait StatisticsFunctionFactory: std::fmt::Debug + Send + Sync {
 
 /// Accumulate statistics.  Must be an associative and commutative over a sequence of `observe` calls.
 pub trait StatisticsFunction: std::fmt::Debug + Send {
-    fn observe(&mut self, hydrated_record: &HydratedMaterializedLogRecord<'_, '_>);
+    fn observe_insert(&mut self, hydrated_record: &HydratedMaterializedLogRecord<'_, '_>);
+    fn observe_delete(&mut self, hydrated_record: &HydratedMaterializedLogRecord<'_, '_>);
     fn output(&self) -> UpdateMetadataValue;
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
 }
 
 #[derive(Debug, Default)]
@@ -45,14 +47,29 @@ pub struct CounterFunction {
     acc: i64,
 }
 
+impl CounterFunction {
+    /// Create a CounterFunction with an initial value.
+    pub fn with_initial_value(value: i64) -> Self {
+        Self { acc: value }
+    }
+}
+
 impl StatisticsFunction for CounterFunction {
-    fn observe(&mut self, _: &HydratedMaterializedLogRecord<'_, '_>) {
+    fn observe_insert(&mut self, _: &HydratedMaterializedLogRecord<'_, '_>) {
         self.acc = self.acc.saturating_add(1);
     }
 
+    fn observe_delete(&mut self, _: &HydratedMaterializedLogRecord<'_, '_>) {
+        self.acc = self.acc.saturating_sub(1);
+    }
+
     fn output(&self) -> UpdateMetadataValue {
         UpdateMetadataValue::Int(self.acc)
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
 }
 
 /// Canonical representation of metadata values tracked by the statistics executor.
@@ -171,31 +188,153 @@ impl Hash for StatisticsValue {
 #[derive(Debug)]
 pub struct StatisticsFunctionExecutor(pub Box<dyn StatisticsFunctionFactory>);
 
+impl StatisticsFunctionExecutor {
+    /// Load existing statistics from the output reader.
+    /// Returns a HashMap with the same structure as the counts HashMap.
+    async fn load_existing_statistics(
+        &self,
+        output_reader: Option<&RecordSegmentReader<'_>>,
+    ) -> Result<
+        HashMap<String, HashMap<StatisticsValue, Box<dyn StatisticsFunction>>>,
+        Box<dyn ChromaError>,
+    > {
+        let mut counts: HashMap<String, HashMap<StatisticsValue, Box<dyn StatisticsFunction>>> =
+            HashMap::default();
+
+        let Some(reader) = output_reader else {
+            return Ok(counts);
+        };
+
+        let max_offset_id = reader.get_max_offset_id();
+        let mut stream = reader.get_data_stream(0..=max_offset_id).await;
+
+        while let Some(record_result) = stream.next().await {
+            let (_, record) = record_result?;
+
+            // Parse the record to extract key, value, type, and count
+            let Some(metadata) = &record.metadata else {
+                continue;
+            };
+
+            let key = match metadata.get("key") {
+                Some(MetadataValue::Str(k)) => k.clone(),
+                _ => continue,
+            };
+
+            let value_type = match metadata.get("type") {
+                Some(MetadataValue::Str(t)) => t.as_str(),
+                _ => continue,
+            };
+
+            let value_str = match metadata.get("value") {
+                Some(MetadataValue::Str(v)) => v.as_str(),
+                _ => continue,
+            };
+
+            let count = match metadata.get("count") {
+                Some(MetadataValue::Int(c)) => *c,
+                _ => continue,
+            };
+
+            // Reconstruct the StatisticsValue from type and value
+            let stats_value = match value_type {
+                "bool" => match value_str {
+                    "true" => StatisticsValue::Bool(true),
+                    "false" => StatisticsValue::Bool(false),
+                    _ => continue,
+                },
+                "int" => match value_str.parse::<i64>() {
+                    Ok(i) => StatisticsValue::Int(i),
+                    _ => continue,
+                },
+                "float" => match value_str.parse::<f64>() {
+                    Ok(f) => StatisticsValue::Float(f),
+                    _ => continue,
+                },
+                "str" => StatisticsValue::Str(value_str.to_string()),
+                "sparse" => match value_str.parse::<u32>() {
+                    Ok(index) => StatisticsValue::SparseVector(index),
+                    _ => continue,
+                },
+                _ => continue,
+            };
+
+            // Create a statistics function initialized with the existing count
+            let stats_function =
+                Box::new(CounterFunction::with_initial_value(count)) as Box<dyn StatisticsFunction>;
+
+            counts
+                .entry(key)
+                .or_default()
+                .insert(stats_value, stats_function);
+        }
+
+        Ok(counts)
+    }
+}
+
 #[async_trait]
 impl AttachedFunctionExecutor for StatisticsFunctionExecutor {
     async fn execute(
         &self,
         input_records: Chunk<HydratedMaterializedLogRecord<'_, '_>>,
         output_reader: Option<&RecordSegmentReader<'_>>,
     ) -> Result<Chunk<LogRecord>, Box<dyn ChromaError>> {
-        let mut counts: HashMap<String, HashMap<StatisticsValue, Box<dyn StatisticsFunction>>> =
-            HashMap::default();
+        // Load existing statistics from output_reader if available
+        let mut counts = self.load_existing_statistics(output_reader).await?;
+
+        // Process new input records and update counts
         for (hydrated_record, _index) in input_records.iter() {
-            // This is only applicable for non-incremental statistics.
-            // TODO(tanujnay112): Change this when we make incremental statistics work.
-            if hydrated_record.get_operation() == MaterializedLogOperation::DeleteExisting {
-                continue;
+            let metadata_delta = hydrated_record.compute_metadata_delta();
+
+            // Decrement counts for deleted metadata
+            for (key, old_value) in metadata_delta.metadata_to_delete {
+                for stats_value in StatisticsValue::from_metadata_value(old_value) {
+                    let inner_map = counts.entry(key.to_string()).or_default();
+                    inner_map
+                        .entry(stats_value)
+                        .or_insert_with(|| self.0.create())
+                        .observe_delete(hydrated_record);
+                }
+            }
+
+            // Decrement counts for old values in updates
+            for (key, (old_value, new_value)) in &metadata_delta.metadata_to_update {
+                for stats_value in StatisticsValue::from_metadata_value(old_value) {
+                    let inner_map = counts.entry(key.to_string()).or_default();
+                    inner_map
+                        .entry(stats_value)
+                        .or_insert_with(|| self.0.create())
+                        .observe_delete(hydrated_record);
+                }
+
+                for stats_value in StatisticsValue::from_metadata_value(new_value) {
+                    let inner_map = counts.entry(key.to_string()).or_default();
+                    inner_map
+                        .entry(stats_value)
+                        .or_insert_with(|| self.0.create())
+                        .observe_insert(hydrated_record);
+                }
             }
 
-            // Use merged_metadata to get the metadata from the hydrated record
-            let metadata = hydrated_record.merged_metadata();
-            for (key, value) in metadata.iter() {
-                let inner_map = counts.entry(key.clone()).or_default();
+            // Increment counts for new values in both updates and inserts
+            for (key, value) in metadata_delta
+                .metadata_to_update
+                .iter()
+                .map(|(k, (_old, new))| (*k, *new))
+                .chain(
+                    metadata_delta
+                        .metadata_to_insert
+                        .iter()
+                        .map(|(k, v)| (*k, *v)),
+                )
+            {
                 for stats_value in StatisticsValue::from_metadata_value(value) {
+                    let inner_map = counts.entry(key.to_string()).or_default();
                     inner_map
                         .entry(stats_value)
                         .or_insert_with(|| self.0.create())
-                        .observe(hydrated_record);
+                        .observe_insert(hydrated_record);
                 }
             }
         }
diff --git a/rust/worker/src/execution/operators/execute_task.rs b/rust/worker/src/execution/operators/execute_task.rs
@@ -16,6 +16,10 @@ use uuid::Uuid;
 use crate::execution::functions::{CounterFunctionFactory, StatisticsFunctionExecutor};
 use crate::execution::operators::materialize_logs::MaterializeLogOutput;
 
+// Constants for CountAttachedFunction
+const COUNT_FUNCTION_OUTPUT_ID: &str = "function_output";
+const COUNT_METADATA_KEY: &str = "total_count";
+
 /// Trait for attached function executors that process input records and produce output records.
 /// Implementors can read from the output collection to maintain state across executions.
 #[async_trait]
@@ -40,31 +44,72 @@ pub trait AttachedFunctionExecutor: Send + Sync + std::fmt::Debug {
 #[derive(Debug)]
 pub struct CountAttachedFunction;
 
+impl CountAttachedFunction {
+    /// Reads the existing count from the output reader.
+    /// Returns 0 if no existing count is found.
+    async fn get_existing_count(output_reader: Option<&RecordSegmentReader<'_>>) -> i64 {
+        let Some(reader) = output_reader else {
+            return 0;
+        };
+
+        // Try to get the existing record with the function output ID
+        let offset_id = match reader
+            .get_offset_id_for_user_id(COUNT_FUNCTION_OUTPUT_ID)
+            .await
+        {
+            Ok(Some(offset_id)) => offset_id,
+            _ => return 0,
+        };
+
+        // Get the data record for this offset id
+        let data_record = match reader.get_data_for_offset_id(offset_id).await {
+            Ok(Some(data_record)) => data_record,
+            _ => return 0,
+        };
+
+        // Extract total_count from metadata
+        if let Some(metadata) = &data_record.metadata {
+            if let Some(chroma_types::MetadataValue::Int(count)) = metadata.get(COUNT_METADATA_KEY)
+            {
+                return *count;
+            }
+        }
+
+        0
+    }
+}
+
 #[async_trait]
 impl AttachedFunctionExecutor for CountAttachedFunction {
     async fn execute(
         &self,
         input_records: Chunk<HydratedMaterializedLogRecord<'_, '_>>,
-        _output_reader: Option<&RecordSegmentReader<'_>>,
+        output_reader: Option<&RecordSegmentReader<'_>>,
     ) -> Result<Chunk<LogRecord>, Box<dyn ChromaError>> {
         let records_count = input_records.len() as i64;
-        let new_total_count = records_count;
+
+        // Read existing count from output_reader if available
+        let existing_count = Self::get_existing_count(output_reader).await;
+        let new_total_count = existing_count + records_count;
 
         // Create output record with updated count
         let mut metadata = std::collections::HashMap::new();
         metadata.insert(
-            "total_count".to_string(),
+            COUNT_METADATA_KEY.to_string(),
             UpdateMetadataValue::Int(new_total_count),
         );
 
         let output_record = LogRecord {
             log_offset: 0,
             record: OperationRecord {
-                id: "function_output".to_string(),
+                id: COUNT_FUNCTION_OUTPUT_ID.to_string(),
                 embedding: Some(vec![0.0]),
                 encoding: None,
                 metadata: Some(metadata),
-                document: Some(format!("Processed {} records", records_count)),
+                document: Some(format!(
+                    "Last processed {} records (total: {})",
+                    records_count, new_total_count
+                )),
                 operation: Operation::Upsert,
             },
         };
@@ -126,6 +171,8 @@ pub struct ExecuteAttachedFunctionInput {
     pub output_record_segment: Segment,
     /// Blockfile provider for reading segments
     pub blockfile_provider: BlockfileProvider,
+
+    pub is_rebuild: bool,
 }
 
 /// Output from the ExecuteAttachedFunction operator
@@ -190,19 +237,23 @@ impl Operator<ExecuteAttachedFunctionInput, ExecuteAttachedFunctionOutput>
         );
 
         // Create record segment reader from the output collection's record segment
-        let record_segment_reader = match Box::pin(RecordSegmentReader::from_segment(
-            &input.output_record_segment,
-            &input.blockfile_provider,
-        ))
-        .await
-        {
-            Ok(reader) => Some(reader),
-            Err(e) if matches!(*e, RecordSegmentReaderCreationError::UninitializedSegment) => {
-                // Output collection has no data yet - this is the first run
-                tracing::info!("[ExecuteAttachedFunction]: Output segment uninitialized - first attached function run");
-                None
+        let record_segment_reader = if input.is_rebuild {
+            None
+        } else {
+            match Box::pin(RecordSegmentReader::from_segment(
+                &input.output_record_segment,
+                &input.blockfile_provider,
+            ))
+            .await
+            {
+                Ok(reader) => Some(reader),
+                Err(e) if matches!(*e, RecordSegmentReaderCreationError::UninitializedSegment) => {
+                    // Output collection has no data yet - this is the first run
+                    tracing::info!("[ExecuteAttachedFunction]: Output segment uninitialized - first attached function run");
+                    None
+                }
+                Err(e) => return Err((*e).into()),
             }
-            Err(e) => return Err((*e).into()),
         };
 
         // Process all materialized logs and hydrate the records
diff --git a/rust/worker/src/execution/orchestration/attached_function_orchestrator.rs b/rust/worker/src/execution/orchestration/attached_function_orchestrator.rs
@@ -715,6 +715,7 @@ impl Handler<TaskResult<CollectionAndSegments, GetCollectionAndSegmentsError>>
             completion_offset: collection_info.pulled_log_offset as u64, // Use the completion offset from input collection
             output_record_segment: message.record_segment.clone(),
             blockfile_provider: self.output_context.blockfile_provider.clone(),
+            is_rebuild: self.output_context.is_rebuild,
         };
 
         let task = wrap(
