Merge branch 'main' into feat/table-features-and-write-tests

Author: zachschuermann

kernel/src/actions/visitors.rs

@@ -2,12 +2,12 @@
 //! [`crate::engine_data::EngineData`] types.
 
 use std::collections::HashMap;
-use std::sync::LazyLock;
+use std::sync::{Arc, LazyLock};
 
 use delta_kernel_derive::internal_api;
 
 use crate::engine_data::{GetData, RowVisitor, TypedGetData as _};
-use crate::schema::{column_name, ColumnName, ColumnNamesAndTypes, DataType, ToSchema as _};
+use crate::schema::{column_name, ColumnName, ColumnNamesAndTypes, DataType, Schema, StructField};
 use crate::utils::require;
 use crate::{DeltaResult, Error};
 
@@ -600,14 +600,82 @@ pub(crate) fn visit_protocol_at<'a>(
     Ok(Some(protocol))
 }
 
+/// This visitor extracts the in-commit timestamp (ICT) from a CommitInfo action in the log it is
+/// present. The [`EngineData`] being visited must have the schema defined in
+/// [`InCommitTimestampVisitor::schema`].
+///
+/// Only the a single row of the engine data is checked (the first row). This is because in-commit
+/// timestamps requires that the CommitInfo containing the ICT be the first action in the log.
+#[allow(unused)]
+#[derive(Default)]
+pub(crate) struct InCommitTimestampVisitor {
+    pub(crate) in_commit_timestamp: Option<i64>,
+}
+
+impl InCommitTimestampVisitor {
+    #[allow(unused)]
+    /// Get the schema that the visitor expects the data to have.
+    pub(crate) fn schema() -> Arc<Schema> {
+        static SCHEMA: LazyLock<Arc<Schema>> = LazyLock::new(|| {
+            let ict_type = StructField::new("inCommitTimestamp", DataType::LONG, true);
+            Arc::new(StructType::new(vec![StructField::new(
+                COMMIT_INFO_NAME,
+                StructType::new([ict_type]),
+                true,
+            )]))
+        });
+        SCHEMA.clone()
+    }
+}
+impl RowVisitor for InCommitTimestampVisitor {
+    fn selected_column_names_and_types(
+        &self,
+    ) -> (&'static [crate::schema::ColumnName], &'static [DataType]) {
+        static NAMES_AND_TYPES: LazyLock<ColumnNamesAndTypes> = LazyLock::new(|| {
+            let names = vec![column_name!("commitInfo.inCommitTimestamp")];
+            let types = vec![DataType::LONG];
+
+            (names, types).into()
+        });
+        NAMES_AND_TYPES.as_ref()
+    }
+
+    fn visit<'a>(
+        &mut self,
+        row_count: usize,
+        getters: &[&'a dyn crate::engine_data::GetData<'a>],
+    ) -> DeltaResult<()> {
+        require!(
+            getters.len() == 1,
+            Error::InternalError(format!(
+                "Wrong number of InCommitTimestampVisitor getters: {}",
+                getters.len()
+            ))
+        );
+
+        // If the batch is empty, return
+        if row_count == 0 {
+            return Ok(());
+        }
+        // CommitInfo must be the first action in a commit
+        if let Some(in_commit_timestamp) = getters[0].get_long(0, "commitInfo.inCommitTimestamp")? {
+            self.in_commit_timestamp = Some(in_commit_timestamp);
+        }
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     use crate::arrow::array::StringArray;
 
+    use crate::engine::sync::SyncEngine;
+    use crate::expressions::{column_expr, Expression};
     use crate::table_features::{ReaderFeature, WriterFeature};
     use crate::utils::test_utils::{action_batch, parse_json_batch};
+    use crate::Engine;
 
     #[test]
     fn test_parse_protocol() -> DeltaResult<()> {
@@ -985,4 +1053,58 @@ mod tests {
             .unwrap();
         assert!(domain_metadata_visitor.domain_metadatas.is_empty());
     }
+
+    /*************************************
+     *  In-commit timestamp visitor tests *
+     **************************************/
+
+    fn add_action() -> &'static str {
+        r#"{"add":{"path":"file1","partitionValues":{"c1":"6","c2":"a"},"size":452,"modificationTime":1670892998137,"dataChange":true}}"#
+    }
+    fn commit_info_action() -> &'static str {
+        r#"{"commitInfo":{"inCommitTimestamp":1677811178585, "timestamp":1677811178585,"operation":"WRITE","operationParameters":{"mode":"ErrorIfExists","partitionBy":"[]"},"isolationLevel":"WriteSerializable","isBlindAppend":true,"operationMetrics":{"numFiles":"1","numOutputRows":"10","numOutputBytes":"635"},"engineInfo":"Databricks-Runtime/<unknown>","txnId":"a6a94671-55ef-450e-9546-b8465b9147de"}}"#
+    }
+
+    fn transform_batch(batch: Box<dyn EngineData>) -> Box<dyn EngineData> {
+        let engine = SyncEngine::new();
+        engine
+            .evaluation_handler()
+            .new_expression_evaluator(
+                get_log_schema().clone(),
+                Expression::Struct(vec![Expression::Struct(vec![column_expr!(
+                    "commitInfo.inCommitTimestamp"
+                )])]),
+                InCommitTimestampVisitor::schema().into(),
+            )
+            .evaluate(batch.as_ref())
+            .unwrap()
+    }
+
+    // Helper function to reduce duplication in tests
+    fn run_timestamp_visitor_test(json_strings: Vec<&str>, expected_timestamp: Option<i64>) {
+        let json_strings: StringArray = json_strings.into();
+        let batch = parse_json_batch(json_strings);
+        let batch = transform_batch(batch);
+        let mut visitor = InCommitTimestampVisitor::default();
+        visitor.visit_rows_of(batch.as_ref()).unwrap();
+        assert_eq!(visitor.in_commit_timestamp, expected_timestamp);
+    }
+
+    #[test]
+    fn commit_info_not_first() {
+        run_timestamp_visitor_test(vec![add_action(), commit_info_action()], None);
+    }
+
+    #[test]
+    fn commit_info_not_present() {
+        run_timestamp_visitor_test(vec![add_action()], None);
+    }
+
+    #[test]
+    fn commit_info_get() {
+        run_timestamp_visitor_test(
+            vec![commit_info_action(), add_action()],
+            Some(1677811178585), // Retrieved ICT
+        );
+    }
 }

kernel/src/history_manager/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod search;