i put in design, cursor help fight rust compiler. high level idea to review.

Author: MannDP

kernel/src/transaction/mod.rs

@@ -202,22 +202,9 @@ impl Transaction {
         let commit_info_action =
             commit_info.into_engine_data(get_log_commit_info_schema().clone(), engine);
 
-        // Step 3: Generate add actions with or without row tracking metadata
+        // Step 3: Generate add actions (handles row tracking internally)
         let commit_version = self.read_snapshot.version() + 1;
-        let add_actions = if self
-            .read_snapshot
-            .table_configuration()
-            .should_write_row_tracking()
-        {
-            self.generate_adds_with_row_tracking(engine, commit_version)?
-        } else {
-            self.generate_adds(
-                engine,
-                self.add_files_metadata.iter().map(|a| Ok(a.deref())),
-                add_files_schema().clone(),
-                as_log_add_schema(with_stats_col(mandatory_add_file_schema())),
-            )
-        };
+        let add_actions = self.generate_adds(engine, commit_version)?;
 
         let domain_metadata_actions = Self::generate_domain_metadata_actions(
             engine,
@@ -368,22 +355,105 @@ impl Transaction {
         self.add_files_metadata.push(add_metadata);
     }
 
-    /// Convert file metadata provided by the engine into protocol-compliant add actions.
-    fn generate_adds<'a, I, T>(
+    /// Generate add actions, handling row tracking internally if needed
+    fn generate_adds<'a>(
+        &'a self,
+        engine: &dyn Engine,
+        commit_version: u64,
+    ) -> DeltaResult<EngineDataResultIterator<'a>> {
+        if self.add_files_metadata.is_empty() {
+            return Ok(Box::new(iter::empty()));
+        }
+
+        let needs_row_tracking = self
+            .read_snapshot
+            .table_configuration()
+            .should_write_row_tracking();
+
+        // Set up variables based on row tracking needs, then make one clean call
+        let (input_schema, output_schema, domain_metadata_action): (
+            SchemaRef,
+            SchemaRef,
+            Option<DeltaResult<Box<dyn EngineData>>>,
+        ) = if needs_row_tracking {
+            let row_id_high_water_mark =
+                RowTrackingDomainMetadata::get_high_water_mark(&self.read_snapshot, engine)?;
+            let mut row_tracking_visitor = RowTrackingVisitor::new(row_id_high_water_mark);
+            let mut base_row_id_batches = Vec::with_capacity(self.add_files_metadata.len());
+
+            // Visit all files to collect row tracking info
+            for add_files_batch in &self.add_files_metadata {
+                row_tracking_visitor.visit_rows_of(add_files_batch.deref())?;
+                base_row_id_batches.push(row_tracking_visitor.base_row_ids.clone());
+            }
+
+            // Create domain metadata action
+            let domain_metadata = DomainMetadata::try_from(RowTrackingDomainMetadata::new(
+                row_tracking_visitor.row_id_high_water_mark,
+            ))?;
+            let domain_metadata_action =
+                domain_metadata.into_engine_data(get_log_domain_metadata_schema().clone(), engine);
+
+            // Extend add files with row tracking columns (mutate in place)
+            for (add_files_batch, base_row_ids) in
+                self.add_files_metadata.iter().zip(base_row_id_batches)
+            {
+                let commit_versions = vec![commit_version as i64; base_row_ids.len()];
+                let base_row_ids_array =
+                    ArrayData::try_new(ArrayType::new(DataType::LONG, true), base_row_ids)?;
+                let commit_versions_array =
+                    ArrayData::try_new(ArrayType::new(DataType::LONG, true), commit_versions)?;
+
+                add_files_batch.append_columns(
+                    with_row_tracking_cols(&Arc::new(StructType::new(vec![]))),
+                    vec![base_row_ids_array, commit_versions_array],
+                )?;
+            }
+
+            (
+                with_row_tracking_cols(add_files_schema()),
+                as_log_add_schema(with_row_tracking_cols(&with_stats_col(
+                    mandatory_add_file_schema(),
+                ))),
+                Some(domain_metadata_action),
+            )
+        } else {
+            (
+                add_files_schema().clone(),
+                as_log_add_schema(with_stats_col(mandatory_add_file_schema())),
+                None,
+            )
+        };
+
+        // Single call to create add actions with the same iterator each time
+        let add_actions = self.create_add_actions(
+            engine,
+            self.add_files_metadata.iter().map(|a| Ok(a.deref())),
+            input_schema,
+            output_schema,
+        );
+
+        // Chain with domain metadata if present
+        Ok(Box::new(
+            add_actions.chain(domain_metadata_action.into_iter()),
+        ))
+    }
+
+    /// Convert file metadata into protocol-compliant add actions
+    fn create_add_actions<'a, I, T>(
         &'a self,
         engine: &dyn Engine,
         add_files_metadata: I,
         input_schema: SchemaRef,
         output_schema: SchemaRef,
-    ) -> EngineDataResultIterator<'a>
+    ) -> impl Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send + 'a
     where
         I: Iterator<Item = DeltaResult<T>> + Send + 'a,
         T: Deref<Target = dyn EngineData> + Send + 'a,
     {
         let evaluation_handler = engine.evaluation_handler();
 
-        Box::new(add_files_metadata.map(move |add_files_batch| {
-            // Convert stats to a JSON string and nest the add action in a top-level struct
+        add_files_metadata.map(move |add_files_batch| {
             let adds_expr = Expression::struct_from([Expression::transform(
                 Transform::new_top_level().with_replaced_field(
                     "stats",
@@ -396,74 +466,7 @@ impl Transaction {
                 output_schema.clone().into(),
             );
             adds_evaluator.evaluate(add_files_batch?.deref())
-        }))
-    }
-
-    /// Extend file metadata provided by the engine with row tracking information and convert them into
-    /// protocol-compliant add actions.
-    fn generate_adds_with_row_tracking<'a>(
-        &'a self,
-        engine: &dyn Engine,
-        commit_version: u64,
-    ) -> DeltaResult<EngineDataResultIterator<'a>> {
-        // Return early if we have nothing to add
-        if self.add_files_metadata.is_empty() {
-            return Ok(Box::new(iter::empty()));
-        }
-
-        // Read the current rowIdHighWaterMark from the snapshot's row tracking domain metadata
-        let row_id_high_water_mark =
-            RowTrackingDomainMetadata::get_high_water_mark(&self.read_snapshot, engine)?;
-
-        // Create a row tracking visitor and visit all files to collect row tracking information
-        let mut row_tracking_visitor = RowTrackingVisitor::new(row_id_high_water_mark);
-        let mut base_row_id_batches = Vec::with_capacity(self.add_files_metadata.len());
-
-        // We visit all files with the row visitor before creating the add action iterator
-        // because we need to know the final row ID high water mark to create the domain metadata action
-        for add_files_batch in &self.add_files_metadata {
-            row_tracking_visitor.visit_rows_of(add_files_batch.deref())?;
-            base_row_id_batches.push(row_tracking_visitor.base_row_ids.clone());
-        }
-
-        // Generate a domain metadata action based on the final high water mark
-        let domain_metadata = DomainMetadata::try_from(RowTrackingDomainMetadata::new(
-            row_tracking_visitor.row_id_high_water_mark,
-        ))?;
-        let domain_metadata_action =
-            domain_metadata.into_engine_data(get_log_domain_metadata_schema().clone(), engine);
-
-        // Create an iterator that pairs each add action with its row tracking metadata
-        let extended_add_files_metadata =
-            self.add_files_metadata.iter().zip(base_row_id_batches).map(
-                move |(add_files_batch, base_row_ids)| {
-                    let commit_versions = vec![commit_version as i64; base_row_ids.len()];
-                    let base_row_ids =
-                        ArrayData::try_new(ArrayType::new(DataType::LONG, true), base_row_ids)?;
-                    let row_commit_versions =
-                        ArrayData::try_new(ArrayType::new(DataType::LONG, true), commit_versions)?;
-
-                    add_files_batch.append_columns(
-                        with_row_tracking_cols(&Arc::new(StructType::new(vec![]))),
-                        vec![base_row_ids, row_commit_versions],
-                    )
-                },
-            );
-
-        // Generate add actions including row tracking metadata
-        let add_actions = self.generate_adds(
-            engine,
-            extended_add_files_metadata,
-            with_row_tracking_cols(add_files_schema()),
-            as_log_add_schema(with_row_tracking_cols(&with_stats_col(
-                mandatory_add_file_schema(),
-            ))),
-        );
-
-        // Return a chained iterator with add and domain metadata actions
-        Ok(Box::new(
-            add_actions.chain(iter::once(domain_metadata_action)),
-        ))
+        })
     }
 }