refactor: Row tracking write cleanup

lbhm · lbhm · commit 52102eb547b5 · 2025-09-18T21:00:21.000+02:00
diff --git a/kernel/src/actions/mod.rs b/kernel/src/actions/mod.rs
@@ -926,14 +926,12 @@ impl DomainMetadata {
 mod tests {
     use super::*;
     use crate::{
-        arrow::{
-            array::{
-                Array, BooleanArray, Int32Array, Int64Array, ListArray, ListBuilder, MapBuilder,
-                MapFieldNames, RecordBatch, StringArray, StringBuilder, StructArray,
-            },
-            datatypes::{DataType as ArrowDataType, Field, Schema},
-            json::ReaderBuilder,
+        arrow::array::{
+            Array, BooleanArray, Int32Array, Int64Array, ListArray, ListBuilder, MapBuilder,
+            MapFieldNames, RecordBatch, StringArray, StringBuilder, StructArray,
         },
+        arrow::datatypes::{DataType as ArrowDataType, Field, Schema},
+        arrow::json::ReaderBuilder,
         engine::{arrow_data::ArrowEngineData, arrow_expression::ArrowEvaluationHandler},
         schema::{ArrayType, DataType, MapType, StructField},
         utils::test_utils::assert_result_error_with_message,
diff --git a/kernel/src/row_tracking.rs b/kernel/src/row_tracking.rs
@@ -81,8 +81,8 @@ pub(crate) struct RowTrackingVisitor {
     /// High water mark for row IDs
     pub(crate) row_id_high_water_mark: i64,
 
-    /// Computed base row IDs of the visited actions
-    pub(crate) base_row_ids: Vec<i64>,
+    /// Computed base row IDs of the visited actions, organized by batch
+    pub(crate) base_row_id_batches: Vec<Vec<i64>>,
 }
 
 impl RowTrackingVisitor {
@@ -94,11 +94,11 @@ impl RowTrackingVisitor {
     /// We verify this hard-coded index in a test.
     const NUM_RECORDS_FIELD_INDEX: usize = 5;
 
-    pub(crate) fn new(row_id_high_water_mark: Option<i64>) -> Self {
+    pub(crate) fn new(row_id_high_water_mark: Option<i64>, num_batches: Option<usize>) -> Self {
         // A table might not have a row ID high water mark yet, so we model the input as an Option<i64>
         Self {
             row_id_high_water_mark: row_id_high_water_mark.unwrap_or(Self::DEFAULT_HIGH_WATER_MARK),
-            base_row_ids: vec![],
+            base_row_id_batches: Vec::with_capacity(num_batches.unwrap_or(0)),
         }
     }
 }
@@ -119,9 +119,8 @@ impl RowVisitor for RowTrackingVisitor {
             ))
         );
 
-        // Reset base row ID vector and allocate the necessary capacity
-        self.base_row_ids.clear();
-        self.base_row_ids.reserve(row_count);
+        // Create a new batch for this visit
+        let mut batch_base_row_ids = Vec::with_capacity(row_count);
 
         let mut current_hwm = self.row_id_high_water_mark;
         for i in 0..row_count {
@@ -133,10 +132,11 @@ impl RowVisitor for RowTrackingVisitor {
                             .to_string(),
                     )
                 })?;
-            self.base_row_ids.push(current_hwm + 1);
+            batch_base_row_ids.push(current_hwm + 1);
             current_hwm += num_records;
         }
 
+        self.base_row_id_batches.push(batch_base_row_ids);
         self.row_id_high_water_mark = current_hwm;
         Ok(())
     }
@@ -205,15 +205,16 @@ mod tests {
 
     #[test]
     fn test_visit_basic_functionality() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(None);
+        let mut visitor = RowTrackingVisitor::new(None, Some(1));
         let num_records_mock = MockGetData::new(vec![Some(10), Some(5), Some(20)]);
         let unit_mock = ();
         let getters = create_getters(&num_records_mock, &unit_mock);
 
         visitor.visit(3, &getters)?;
 
         // Check that base row IDs are calculated correctly
-        assert_eq!(visitor.base_row_ids, vec![0, 10, 15]);
+        assert_eq!(visitor.base_row_id_batches.len(), 1);
+        assert_eq!(visitor.base_row_id_batches[0], vec![0, 10, 15]);
 
         // Check that high water mark is updated correctly
         assert_eq!(visitor.row_id_high_water_mark, 34); // -1 + 10 + 5 + 20
@@ -223,15 +224,16 @@ mod tests {
 
     #[test]
     fn test_visit_with_negative_high_water_mark() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(Some(-5));
+        let mut visitor = RowTrackingVisitor::new(Some(-5), Some(1));
         let num_records_mock = MockGetData::new(vec![Some(3), Some(2)]);
         let unit_mock = ();
         let getters = create_getters(&num_records_mock, &unit_mock);
 
         visitor.visit(2, &getters)?;
 
         // Base row IDs should start from high_water_mark + 1
-        assert_eq!(visitor.base_row_ids, vec![-4, -1]); // -5+1=-4, then -4+3=-1
+        assert_eq!(visitor.base_row_id_batches.len(), 1);
+        assert_eq!(visitor.base_row_id_batches[0], vec![-4, -1]); // -5+1=-4, then -4+3=-1
 
         // High water mark should be updated
         assert_eq!(visitor.row_id_high_water_mark, 0); // -5 + 3 + 2 = 0
@@ -241,15 +243,16 @@ mod tests {
 
     #[test]
     fn test_visit_with_zero_records() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(Some(10));
+        let mut visitor = RowTrackingVisitor::new(Some(10), Some(1));
         let num_records_mock = MockGetData::new(vec![Some(0), Some(0), Some(5)]);
         let unit_mock = ();
         let getters = create_getters(&num_records_mock, &unit_mock);
 
         visitor.visit(3, &getters)?;
 
         // Base row IDs should still be assigned even for zero-record files
-        assert_eq!(visitor.base_row_ids, vec![11, 11, 11]);
+        assert_eq!(visitor.base_row_id_batches.len(), 1);
+        assert_eq!(visitor.base_row_id_batches[0], vec![11, 11, 11]);
 
         // High water mark should only increase by non-zero records
         assert_eq!(visitor.row_id_high_water_mark, 15); // 10 + 0 + 0 + 5
@@ -259,23 +262,54 @@ mod tests {
 
     #[test]
     fn test_visit_empty_batch() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(Some(42));
+        let mut visitor = RowTrackingVisitor::new(Some(42), None);
         let num_records_mock = MockGetData::new(vec![]);
         let unit_mock = ();
         let getters = create_getters(&num_records_mock, &unit_mock);
 
         visitor.visit(0, &getters)?;
 
         // Should handle empty batch gracefully
-        assert!(visitor.base_row_ids.is_empty());
+        assert_eq!(visitor.base_row_id_batches.len(), 1);
+        assert!(visitor.base_row_id_batches[0].is_empty());
         assert_eq!(visitor.row_id_high_water_mark, 42); // Should remain unchanged
 
         Ok(())
     }
 
+    #[test]
+    fn test_visit_multiple_batches() -> DeltaResult<()> {
+        let mut visitor = RowTrackingVisitor::new(Some(0), Some(2));
+        let unit_mock = ();
+
+        // First batch
+        let num_records_mock1 = MockGetData::new(vec![Some(10), Some(5)]);
+        let getters1 = create_getters(&num_records_mock1, &unit_mock);
+        visitor.visit(2, &getters1)?;
+
+        // Second batch
+        let num_records_mock2 = MockGetData::new(vec![Some(3), Some(7), Some(2)]);
+        let getters2 = create_getters(&num_records_mock2, &unit_mock);
+        visitor.visit(3, &getters2)?;
+
+        // Check that we have two batches
+        assert_eq!(visitor.base_row_id_batches.len(), 2);
+
+        // Check first batch: starts at 1, then 11
+        assert_eq!(visitor.base_row_id_batches[0], vec![1, 11]);
+
+        // Check second batch: starts at 16, then 19, then 26
+        assert_eq!(visitor.base_row_id_batches[1], vec![16, 19, 26]);
+
+        // Check final high water mark: 0 + 10 + 5 + 3 + 7 + 2 = 27
+        assert_eq!(visitor.row_id_high_water_mark, 27);
+
+        Ok(())
+    }
+
     #[test]
     fn test_visit_wrong_getter_count() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(Some(0));
+        let mut visitor = RowTrackingVisitor::new(Some(0), None);
         let unit_mock = ();
         let wrong_getters: Vec<&dyn GetData<'_>> = vec![&unit_mock]; // Only one getter instead of expected count
 
@@ -287,7 +321,7 @@ mod tests {
 
     #[test]
     fn test_visit_missing_num_records() -> DeltaResult<()> {
-        let mut visitor = RowTrackingVisitor::new(Some(0));
+        let mut visitor = RowTrackingVisitor::new(Some(0), None);
         let num_records_mock = MockGetData::new(vec![None]); // Missing numRecords
         let unit_mock = ();
         let getters = create_getters(&num_records_mock, &unit_mock);
@@ -303,7 +337,7 @@ mod tests {
 
     #[test]
     fn test_selected_column_names_and_types() {
-        let visitor = RowTrackingVisitor::new(Some(0));
+        let visitor = RowTrackingVisitor::new(Some(0), None);
         let (names, types) = visitor.selected_column_names_and_types();
 
         // Should return the same as add_files_schema().leaves(None)
diff --git a/kernel/src/transaction/mod.rs b/kernel/src/transaction/mod.rs
@@ -221,8 +221,8 @@ impl Transaction {
         let commit_path =
             ParsedLogPath::new_commit(self.read_snapshot.table_root(), commit_version)?;
         let actions = iter::once(commit_info_action)
-            .chain(set_transaction_actions)
-            .chain(add_actions);
+            .chain(add_actions)
+            .chain(set_transaction_actions);
 
         let json_handler = engine.json_handler();
         match json_handler.write_json_file(&commit_path.location, Box::new(actions), false) {
@@ -357,28 +357,35 @@ impl Transaction {
             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());
+        let mut row_tracking_visitor =
+            RowTrackingVisitor::new(row_id_high_water_mark, Some(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());
         }
 
+        // Deconstruct the row tracking visitor to avoid borrowing issues
+        let RowTrackingVisitor {
+            base_row_id_batches,
+            row_id_high_water_mark,
+        } = row_tracking_visitor;
+
         // 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 =
+            DomainMetadata::try_from(RowTrackingDomainMetadata::new(row_id_high_water_mark))?;
         let domain_metadata_action =
             domain_metadata.into_engine_data(get_log_domain_metadata_schema().clone(), engine);
 
+        let commit_version = i64::try_from(commit_version)
+            .map_err(|_| Error::generic("Commit version is too large to fit in i64"))?;
+
         // 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 commit_versions = vec![commit_version; base_row_ids.len()];
                     let base_row_ids =
                         ArrayData::try_new(ArrayType::new(DataType::LONG, true), base_row_ids)?;
                     let row_commit_versions =
