feat: Add write support for row tracking

Author: lbhm

ffi/src/transaction/mod.rs

@@ -130,13 +130,15 @@ mod tests {
     use delta_kernel::schema::{DataType, StructField, StructType};
 
     use delta_kernel::arrow::array::{Array, ArrayRef, Int32Array, StringArray, StructArray};
-    use delta_kernel::arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
+    use delta_kernel::arrow::datatypes::Schema as ArrowSchema;
     use delta_kernel::arrow::ffi::to_ffi;
     use delta_kernel::arrow::json::reader::ReaderBuilder;
     use delta_kernel::arrow::record_batch::RecordBatch;
+    use delta_kernel::engine::arrow_conversion::TryIntoArrow as _;
     use delta_kernel::engine::arrow_data::ArrowEngineData;
     use delta_kernel::parquet::arrow::arrow_writer::ArrowWriter;
     use delta_kernel::parquet::file::properties::WriterProperties;
+    use delta_kernel::transaction::add_files_schema;
 
     use delta_kernel_ffi::engine_data::get_engine_data;
     use delta_kernel_ffi::engine_data::ArrowFFIData;
@@ -183,39 +185,15 @@ mod tests {
         path: &str,
         num_rows: i64,
     ) -> Result<ArrowFFIData, Box<dyn std::error::Error>> {
-        let schema = ArrowSchema::new(vec![
-            Field::new("path", ArrowDataType::Utf8, false),
-            Field::new(
-                "partitionValues",
-                ArrowDataType::Map(
-                    Arc::new(Field::new(
-                        "entries",
-                        ArrowDataType::Struct(
-                            vec![
-                                Field::new("key", ArrowDataType::Utf8, false),
-                                Field::new("value", ArrowDataType::Utf8, true),
-                            ]
-                            .into(),
-                        ),
-                        false,
-                    )),
-                    false,
-                ),
-                false,
-            ),
-            Field::new("size", ArrowDataType::Int64, false),
-            Field::new("modificationTime", ArrowDataType::Int64, false),
-            Field::new("dataChange", ArrowDataType::Boolean, false),
-            Field::new("numRecords", ArrowDataType::Int64, true),
-        ]);
+        let schema: ArrowSchema = add_files_schema().as_ref().try_into_arrow()?;
 
         let current_time: i64 = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap()
             .as_millis() as i64;
 
         let file_metadata = format!(
-            r#"{{"path":"{path}", "partitionValues": {{}}, "size": {num_rows}, "modificationTime": {current_time}, "dataChange": true, "numRecords": {num_rows}}}"#,
+            r#"{{"path":"{path}", "partitionValues": {{}}, "size": {num_rows}, "modificationTime": {current_time}, "dataChange": true, "stats": {{"numRecords": {num_rows}}}}}"#,
         );
 
         create_arrow_ffi_from_json(schema, file_metadata.as_str())

kernel/src/actions/mod.rs

@@ -132,6 +132,14 @@ pub(crate) fn get_log_domain_metadata_schema() -> &'static SchemaRef {
     &LOG_DOMAIN_METADATA_SCHEMA
 }
 
+/// Nest an existing add action schema in an additional [`ADD_NAME`] struct.
+///
+/// This is useful for JSON conversion, as it allows us to wrap a dynamically maintained add action
+/// schema in a top-level "add" struct.
+pub(crate) fn as_log_add_schema(schema: SchemaRef) -> SchemaRef {
+    Arc::new(StructType::new([StructField::nullable(ADD_NAME, schema)]))
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, ToSchema)]
 #[cfg_attr(
     any(test, feature = "internal-api"),
@@ -675,7 +683,7 @@ pub(crate) struct Add {
 
     /// Default generated Row ID of the first row in the file. The default generated Row IDs
     /// of the other rows in the file can be reconstructed by adding the physical index of the
-    /// row within the file to the base Row ID
+    /// row within the file to the base Row ID.
     #[cfg_attr(test, serde(skip_serializing_if = "Option::is_none"))]
     pub base_row_id: Option<i64>,
 
@@ -883,6 +891,17 @@ pub(crate) struct DomainMetadata {
 }
 
 impl DomainMetadata {
+    /// Create a new DomainMetadata action.
+    // TODO: Discuss if we should remove `removed` from this method and introduce a dedicated
+    // method for removed domain metadata.
+    pub(crate) fn new(domain: String, configuration: String, removed: bool) -> Self {
+        DomainMetadata {
+            domain,
+            configuration,
+            removed,
+        }
+    }
+
     // returns true if the domain metadata is an system-controlled domain (all domains that start
     // with "delta.")
     #[allow(unused)]
@@ -904,7 +923,6 @@ mod tests {
         engine::arrow_data::ArrowEngineData,
         engine::arrow_expression::ArrowEvaluationHandler,
         schema::{ArrayType, DataType, MapType, StructField},
-        utils::test_utils::assert_result_error_with_message,
         Engine, EvaluationHandler, JsonHandler, ParquetHandler, StorageHandler,
     };
     use serde_json::json;
@@ -1298,22 +1316,11 @@ mod tests {
                 WriterFeature::AppendOnly,
                 WriterFeature::DeletionVectors,
                 WriterFeature::Invariants,
+                WriterFeature::RowTracking,
             ]),
         )
         .unwrap();
         assert!(protocol.ensure_write_supported().is_ok());
-
-        let protocol = Protocol::try_new(
-            3,
-            7,
-            Some([ReaderFeature::DeletionVectors]),
-            Some([WriterFeature::RowTracking]),
-        )
-        .unwrap();
-        assert_result_error_with_message(
-            protocol.ensure_write_supported(),
-            r#"Unsupported: Unknown WriterFeatures: "rowTracking". Supported WriterFeatures: "appendOnly", "deletionVectors", "invariants", "timestampNtz", "variantType", "variantType-preview", "variantShredding-preview""#,
-        );
     }
 
     #[test]

kernel/src/engine/default/parquet.rs

@@ -5,7 +5,8 @@ use std::ops::Range;
 use std::sync::Arc;
 
 use crate::arrow::array::builder::{MapBuilder, MapFieldNames, StringBuilder};
-use crate::arrow::array::{BooleanArray, Int64Array, RecordBatch, StringArray};
+use crate::arrow::array::{BooleanArray, Int64Array, RecordBatch, StringArray, StructArray};
+use crate::arrow::datatypes::{DataType, Field};
 use crate::parquet::arrow::arrow_reader::{
     ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder,
 };
@@ -24,6 +25,7 @@ use crate::engine::arrow_utils::{fixup_parquet_read, generate_mask, get_requeste
 use crate::engine::default::executor::TaskExecutor;
 use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
 use crate::schema::SchemaRef;
+use crate::transaction::add_files_schema;
 use crate::{
     DeltaResult, EngineData, Error, FileDataReadResultIterator, FileMeta, ParquetHandler,
     PredicateRef,
@@ -55,7 +57,10 @@ impl DataFileMetadata {
         }
     }
 
-    // convert DataFileMetadata into a record batch which matches the 'add_files_schema' schema
+    /// Convert DataFileMetadata into a record batch which matches the schema returned by
+    /// [`add_files_schema`].
+    ///
+    /// [`add_files_schema`]: crate::transaction::add_files_schema
     fn as_record_batch(
         &self,
         partition_values: &HashMap<String, String>,
@@ -70,8 +75,6 @@ impl DataFileMetadata {
                 },
             num_records,
         } = self;
-        let add_files_schema = crate::transaction::add_files_schema();
-
         // create the record batch of the write metadata
         let path = Arc::new(StringArray::from(vec![location.to_string()]));
         let key_builder = StringBuilder::new();
@@ -95,17 +98,22 @@ impl DataFileMetadata {
         let size = Arc::new(Int64Array::from(vec![size]));
         let data_change = Arc::new(BooleanArray::from(vec![data_change]));
         let modification_time = Arc::new(Int64Array::from(vec![*last_modified]));
-        let num_records = Arc::new(Int64Array::from(vec![*num_records as i64]));
+        let stats = Arc::new(StructArray::try_new_with_length(
+            vec![Field::new("numRecords", DataType::Int64, true)].into(),
+            vec![Arc::new(Int64Array::from(vec![*num_records as i64]))],
+            None,
+            1,
+        )?);
 
         Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
-            Arc::new(add_files_schema.as_ref().try_into_arrow()?),
+            Arc::new(add_files_schema().as_ref().try_into_arrow()?),
             vec![
                 path,
                 partitions,
                 size,
                 modification_time,
                 data_change,
-                num_records,
+                stats,
             ],
         )?)))
     }
@@ -502,6 +510,14 @@ mod tests {
         partition_values_builder.values().append_value("a");
         partition_values_builder.append(true).unwrap();
         let partition_values = partition_values_builder.finish();
+        let stats_struct = StructArray::try_new_with_length(
+            vec![Field::new("numRecords", DataType::Int64, true)].into(),
+            vec![Arc::new(Int64Array::from(vec![num_records as i64]))],
+            None,
+            1,
+        )
+        .unwrap();
+
         let expected = RecordBatch::try_new(
             schema,
             vec![
@@ -510,7 +526,7 @@ mod tests {
                 Arc::new(Int64Array::from(vec![size as i64])),
                 Arc::new(Int64Array::from(vec![last_modified])),
                 Arc::new(BooleanArray::from(vec![data_change])),
-                Arc::new(Int64Array::from(vec![num_records as i64])),
+                Arc::new(stats_struct),
             ],
         )
         .unwrap();

kernel/src/lib.rs

@@ -89,6 +89,7 @@ pub mod checkpoint;
 pub mod engine_data;
 pub mod error;
 pub mod expressions;
+pub mod row_tracking;
 pub mod scan;
 pub mod schema;
 pub mod snapshot;

kernel/src/row_tracking.rs

@@ -0,0 +1,94 @@
+use std::sync::LazyLock;
+
+use serde::{Deserialize, Serialize};
+
+use delta_kernel_derive::{internal_api, ToSchema};
+
+use crate::engine_data::{GetData, RowVisitor, TypedGetData as _};
+use crate::schema::{ColumnName, ColumnNamesAndTypes, DataType};
+use crate::transaction::add_files_schema;
+use crate::utils::require;
+use crate::{DeltaResult, Error};
+
+pub(crate) const ROW_TRACKING_DOMAIN_NAME: &str = "delta.rowTracking";
+
+#[derive(Debug, Clone, PartialEq, Eq, ToSchema, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct RowTrackingDomainMetadata {
+    // The Delta spec does not define unsigned integers, so we use i64 here
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) row_id_high_water_mark: Option<i64>,
+}
+
+/// A row visitor that iterates over preliminary [`Add`] actions as returned by the engine and
+/// computes a base row ID for each action.
+/// It expects to visit engine data conforming to the schema returned by [`add_files_schema()`].
+///
+/// This visitor it only required for the row tracking write path. The read path will be completely
+/// implemented via expressions.
+///
+/// [`Add`]: delta_kernel::actions::Add
+#[internal_api]
+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>,
+}
+
+impl RowTrackingVisitor {
+    /// Field index for "numRecords" in [`add_files_schema()`]
+    const NUM_RECORDS_FIELD_INDEX: usize = 5;
+
+    /// Default value for an absent high water mark
+    const DEFAULT_HWM: i64 = -1;
+
+    #[internal_api]
+    pub(crate) fn new(row_id_high_water_mark: Option<i64>) -> 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_HWM),
+            base_row_ids: vec![],
+        }
+    }
+}
+
+impl RowVisitor for RowTrackingVisitor {
+    fn selected_column_names_and_types(&self) -> (&'static [ColumnName], &'static [DataType]) {
+        static NAMES_AND_TYPES: LazyLock<ColumnNamesAndTypes> =
+            LazyLock::new(|| add_files_schema().leaves(None));
+        NAMES_AND_TYPES.as_ref()
+    }
+
+    fn visit<'a>(&mut self, row_count: usize, getters: &[&'a dyn GetData<'a>]) -> DeltaResult<()> {
+        require!(
+            getters.len() == add_files_schema().fields_len(),
+            Error::InternalError(format!(
+                "Wrong number of RowTrackingVisitor getters: {}",
+                getters.len()
+            ))
+        );
+
+        // Reset base row ID vector and allocate the necessary capacity
+        self.base_row_ids.clear();
+        self.base_row_ids.reserve(row_count);
+
+        let mut current_hwm = self.row_id_high_water_mark;
+        for i in 0..row_count {
+            let num_records: i64 = getters[Self::NUM_RECORDS_FIELD_INDEX]
+                .get_opt(i, "numRecords")?
+                .ok_or_else(|| {
+                    Error::InternalError(
+                        "numRecords must be present in Add actions when row tracking is enabled."
+                            .to_string(),
+                    )
+                })?;
+            self.base_row_ids.push(current_hwm + 1);
+            current_hwm += num_records;
+        }
+
+        self.row_id_high_water_mark = current_hwm;
+        Ok(())
+    }
+}

kernel/src/table_configuration.rs

@@ -342,6 +342,61 @@ impl TableConfiguration {
             )),
         }
     }
+
+    /// Returns `true` if the table supports writing domain metadata.
+    ///
+    /// To support this feature the table must:
+    /// - Have a min_writer_version of 7.
+    /// - Have the [`WriterFeature::DomainMetadata`] writer feature.
+    #[allow(unused)]
+    pub(crate) fn is_domain_metadata_supported(&self) -> bool {
+        self.protocol().min_writer_version() == 7
+            && self
+                .protocol()
+                .has_writer_feature(&WriterFeature::DomainMetadata)
+    }
+
+    /// Returns `true` if the table supports writing row tracking metadata.
+    ///
+    /// To support this feature the table must:
+    /// - Have a min_writer_version of 7.
+    /// - Have the [`WriterFeature::RowTracking`] writer feature.
+    #[allow(unused)]
+    pub(crate) fn is_row_tracking_supported(&self) -> bool {
+        self.protocol().min_writer_version() == 7
+            && self
+                .protocol()
+                .has_writer_feature(&WriterFeature::RowTracking)
+    }
+
+    /// Returns `true` if row tracking is enabled for this table.
+    ///
+    /// In order to enable row tracking the table must:
+    /// - Support row tracking (see [`Self::is_row_tracking_supported`]).
+    /// - Have the `delta.enableRowTracking` table property set to `true`.
+    #[allow(unused)]
+    pub(crate) fn is_row_tracking_enabled(&self) -> bool {
+        self.is_row_tracking_supported()
+            && self.table_properties().enable_row_tracking.unwrap_or(false)
+    }
+
+    /// Returns `true` if row tracking is suspended for this table.
+    ///
+    /// Row tracking is suspended when the `delta.rowTrackingSuspended` table property is set to `true`.
+    /// Note that:
+    /// - If row tracking is not supported, it is implicitly suspended.
+    /// - Row tracking can be _supported_ but _suspended_ at the same time.
+    /// - Row tracking cannot be _enabled_ while _suspended_.
+    #[allow(unused)]
+    pub(crate) fn is_row_tracking_suspended(&self) -> bool {
+        // We do not check whether row tracking is enabled here but rely on the table properties to
+        // be in a valid state.
+        !self.is_row_tracking_supported()
+            || self
+                .table_properties()
+                .row_tracking_suspended
+                .unwrap_or(false)
+    }
 }
 
 #[cfg(test)]