sicheng comments

Author: tanujnay112

rust/sysdb/src/sysdb.rs

@@ -628,7 +628,6 @@ impl SysDb {
         }
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub async fn flush_compaction_and_attached_function(
         &mut self,
         collections: Vec<CollectionFlushInfo>,

rust/types/src/flush.rs

@@ -158,27 +158,6 @@ impl FlushCompactionResponse {
     }
 }
 
-impl TryFrom<FlushCollectionCompactionAndAttachedFunctionResponse> for FlushCompactionResponse {
-    type Error = FlushCompactionResponseConversionError;
-
-    fn try_from(
-        value: FlushCollectionCompactionAndAttachedFunctionResponse,
-    ) -> Result<Self, Self::Error> {
-        // Use first collection for backward compatibility
-        let first_collection = value
-            .collections
-            .first()
-            .ok_or(FlushCompactionResponseConversionError::MissingCollections)?;
-        let id = Uuid::parse_str(&first_collection.collection_id)
-            .map_err(|_| FlushCompactionResponseConversionError::InvalidUuid)?;
-        Ok(FlushCompactionResponse {
-            collection_id: CollectionUuid(id),
-            collection_version: first_collection.collection_version,
-            last_compaction_time: first_collection.last_compaction_time,
-        })
-    }
-}
-
 impl TryFrom<FlushCollectionCompactionAndAttachedFunctionResponse>
     for FlushCompactionAndAttachedFunctionResponse
 {

rust/worker/src/execution/functions/statistics.rs

@@ -181,7 +181,8 @@ impl AttachedFunctionExecutor for StatisticsFunctionExecutor {
         let mut counts: HashMap<String, HashMap<StatisticsValue, Box<dyn StatisticsFunction>>> =
             HashMap::default();
         for (hydrated_record, _index) in input_records.iter() {
-            // Skip delete operations - they should not be counted in statistics
+            // 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;
             }

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

@@ -50,8 +50,6 @@ impl AttachedFunctionExecutor for CountAttachedFunction {
         let records_count = input_records.len() as i64;
         let new_total_count = records_count;
 
-        println!("new_total_count is {}", new_total_count);
-
         // Create output record with updated count
         let mut metadata = std::collections::HashMap::new();
         metadata.insert(

rust/worker/src/execution/orchestration/apply_logs_orchestrator.rs

@@ -206,7 +206,7 @@ impl ApplyLogsOrchestrator {
         let mut tasks_to_run = Vec::new();
         self.num_materialized_logs += materialized_logs.len() as u64;
 
-        let writers = self.context.get_output_segment_writers()?;
+        let writers = self.context.get_segment_writers()?;
 
         {
             self.num_uncompleted_tasks_by_segment
@@ -255,7 +255,7 @@ impl ApplyLogsOrchestrator {
                 materialized_logs.clone(),
                 writers.record_reader.clone(),
                 self.context
-                    .get_output_collection_info()?
+                    .get_collection_info()?
                     .collection
                     .schema
                     .clone(),
@@ -356,7 +356,7 @@ impl ApplyLogsOrchestrator {
             .add(self.num_materialized_logs, &[]);
 
         self.state = ExecutionState::Register;
-        let collection_info = match self.context.get_output_collection_info() {
+        let collection_info = match self.context.get_collection_info() {
             Ok(collection_info) => collection_info,
             Err(err) => {
                 self.terminate_with_result(Err(err.into()), ctx).await;
@@ -525,7 +525,7 @@ impl Handler<TaskResult<ApplyLogToSegmentWriterOutput, ApplyLogToSegmentWriterOp
 
         if message.segment_type == "MetadataSegmentWriter" {
             if let Some(update) = message.schema_update {
-                let collection_info = match self.context.get_output_collection_info_mut() {
+                let collection_info = match self.context.get_collection_info_mut() {
                     Ok(info) => info,
                     Err(err) => {
                         return self.terminate_with_result(Err(err.into()), ctx).await;
@@ -582,9 +582,7 @@ impl Handler<TaskResult<ApplyLogToSegmentWriterOutput, ApplyLogToSegmentWriterOp
         };
 
         if num_tasks_left == 0 && self.num_uncompleted_materialization_tasks == 0 {
-            let segment_writer = self
-                .context
-                .get_output_segment_writer_by_id(message.segment_id);
+            let segment_writer = self.context.get_segment_writer_by_id(message.segment_id);
             let segment_writer = match self.ok_or_terminate(segment_writer, ctx).await {
                 Some(writer) => writer,
                 None => return,
@@ -614,7 +612,7 @@ impl Handler<TaskResult<CommitSegmentWriterOutput, CommitSegmentWriterOperatorEr
 
         // If the flusher received is a record segment flusher, get the number of keys for the blockfile and set it on the orchestrator
         if let ChromaSegmentFlusher::RecordSegment(record_segment_flusher) = &message.flusher {
-            let collection_info = match self.context.get_output_collection_info_mut() {
+            let collection_info = match self.context.get_collection_info_mut() {
                 Ok(info) => info,
                 Err(err) => {
                     self.terminate_with_result(Err(err.into()), ctx).await;

rust/worker/src/execution/orchestration/attached_function_orchestrator.rs

@@ -230,7 +230,7 @@ impl AttachedFunctionOrchestrator {
         &self,
     ) -> Result<&CollectionCompactInfo, AttachedFunctionOrchestratorError> {
         self.output_context
-            .get_output_collection_info()
+            .get_collection_info()
             .map_err(AttachedFunctionOrchestratorError::CompactionContext)
     }
 
@@ -239,7 +239,7 @@ impl AttachedFunctionOrchestrator {
         &self,
     ) -> Result<CollectionUuid, AttachedFunctionOrchestratorError> {
         self.output_context
-            .get_output_collection_info()
+            .get_collection_info()
             .map(|info| info.collection_id)
             .map_err(AttachedFunctionOrchestratorError::CompactionContext)
     }
@@ -249,9 +249,7 @@ impl AttachedFunctionOrchestrator {
         &mut self,
         collection_info: CollectionCompactInfo,
     ) -> Result<(), CollectionCompactInfo> {
-        self.output_context
-            .output_collection_info
-            .set(collection_info)
+        self.output_context.collection_info.set(collection_info)
     }
 
     /// Get the function context if it has been set
@@ -339,7 +337,7 @@ impl AttachedFunctionOrchestrator {
         // NOTE: We allow writers to be uninitialized for the case when the materialized logs are empty
         let record_reader = self
             .output_context
-            .get_output_segment_writers()
+            .get_segment_writers()
             .ok()
             .and_then(|writers| writers.record_reader);
 
@@ -363,7 +361,7 @@ impl AttachedFunctionOrchestrator {
                 }
             };
 
-            let collection_info = match self.output_context.get_output_collection_info_mut() {
+            let collection_info = match self.output_context.get_collection_info_mut() {
                 Ok(info) => info,
                 Err(err) => {
                     return self.terminate_with_result(Err(err.into()), ctx).await;
@@ -490,7 +488,7 @@ impl Handler<TaskResult<GetAttachedFunctionOutput, GetAttachedFunctionOperatorEr
                 );
 
                 // TODO(tanujnay112): Handle error
-                let _ = self.function_context.set(FunctionContext {
+                let _ = self.set_function_context(FunctionContext {
                     attached_function_id: attached_function.id,
                     function_id: attached_function.function_id,
                     updated_completion_offset: attached_function.completion_offset,