test: cover timezone-aware parsed partitions

Author: DrakeLin

kernel/src/engine/arrow_expression/evaluate_expression.rs

@@ -425,7 +425,7 @@ pub fn evaluate_expression(
                 let values = strings
                     .iter()
                     .map(|raw| raw.and_then(|raw| timezone.parse_timestamp(raw).ok()));
-                Arc::new(TimestampMicrosecondArray::from_iter(values).with_timezone_utc())
+                Arc::new(TimestampMicrosecondArray::from_iter(values).with_timezone("UTC"))
                     as ArrayRef
             } else if can_cast_types(input.data_type(), &target) {
                 cast(&input, &target)?
@@ -3254,6 +3254,10 @@ mod tests {
             None => Expr::cast(col!("s"), target.clone()),
         };
         let result = evaluate_expression(&expr, &batch, Some(&target))?;
+        assert_eq!(
+            result.data_type(),
+            &ArrowDataType::try_from_kernel(&target)?
+        );
         let timestamps = result
             .as_any()
             .downcast_ref::<TimestampMicrosecondArray>()

kernel/src/scan/scan_plan/tests.rs

@@ -15,7 +15,7 @@ use crate::arrow::util::pretty::pretty_format_batches;
 use crate::engine::arrow_data::EngineDataArrowExt as _;
 use crate::engine::sync::SyncEngine;
 use crate::engine::test_delegating::DelegatingEngine;
-use crate::expressions::{col, column_name, lit, Predicate as Pred};
+use crate::expressions::{col, column_name, lit, Predicate as Pred, Scalar};
 use crate::plans::ir::nodes::Operator;
 use crate::plans::Operation as PlanOperation;
 use crate::scan::{PartitionValuesOptions, Scan, StatsOptions, StructStats};
@@ -253,6 +253,58 @@ fn declarative_metadata_matches_imperative_scan(
     }
 }
 
+#[rstest]
+#[case::json_commit(LogState::with_latest_version(1), FeatureSet::new())]
+#[case::v1_checkpoint(
+    LogState::with_latest_version(1).with_checkpoint_at([1]),
+    FeatureSet::new()
+)]
+#[case::v2_sidecar(
+    LogState::with_latest_version(1)
+        .with_checkpoint_at([1])
+        .with_sidecars_if_enabled(None),
+    FeatureSet::new().v2_checkpoint()
+)]
+#[case::column_mapping(
+    LogState::with_latest_version(1).with_checkpoint_at([1]),
+    FeatureSet::new().column_mapping("id")
+)]
+fn declarative_metadata_matches_imperative_with_reader_timezone(
+    #[case] log_state: LogState,
+    #[case] features: FeatureSet,
+) -> DeltaResult<()> {
+    let table = TestTableBuilder::new()
+        .with_log_state(log_state)
+        .with_features(features)
+        .with_table_config(TableConfig::new().write_stats_as_struct(true))
+        .with_data_layout(DataLayoutConfig::PartitionedAllTypes)
+        .build()
+        .expect("build reader-timezone parity table");
+    let engine = SyncEngine::new_with_store(table.store().clone());
+    let snapshot = Snapshot::builder_for(table.table_root()).build(&engine)?;
+    let partition_values =
+        || PartitionValuesOptions::with_struct().with_timestamp_timezone("America/Los_Angeles");
+    let predicate: PredicateRef = Pred::gt(col!("part_ts"), Scalar::Timestamp(0)).into();
+
+    let expected = imperative_metadata(
+        snapshot
+            .clone()
+            .scan_builder()
+            .with_partition_values(partition_values())
+            .with_predicate(predicate.clone())
+            .build()?,
+        &engine,
+    )?;
+    let scan = snapshot
+        .scan_builder()
+        .with_partition_values(partition_values())
+        .with_predicate(predicate)
+        .build()?;
+    let actual = declarative_metadata(&scan, &engine)?;
+
+    assert_metadata_eq(&actual, &expected, "reader-timezone metadata parity")
+}
+
 #[rstest]
 #[case::parquet_manifest("v2-checkpoints-parquet-with-sidecars")]
 #[case::json_manifest("v2-checkpoints-json-with-sidecars")]

kernel/tests/integration/parsed_partition_values.rs

@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 use delta_kernel::arrow::array::{
     Array, BinaryArray, BooleanArray, Int32Array, RecordBatch, StringArray, StructArray,
+    TimestampMicrosecondArray,
 };
 use delta_kernel::arrow::compute::filter_record_batch;
 use delta_kernel::engine::arrow_data::ArrowEngineData;
@@ -391,3 +392,194 @@ async fn empty_string_partition_pruning(#[values(false, true)] native_checkpoint
         "empty-bytes file must be pruned under p_bin = X'6f74686572'"
     );
 }
+
+async fn write_timezone_partition_table(table_path: &std::path::Path, timestamp: &str) -> Url {
+    std::fs::create_dir_all(table_path).unwrap();
+    let url = Url::from_directory_path(table_path).unwrap();
+    let table_root = url.to_string();
+    let store: Arc<DynObjectStore> = Arc::new(LocalFileSystem::new());
+    let schema_string = serde_json::json!({
+        "type": "struct",
+        "fields": [
+            {"name": "p_ts", "type": "timestamp", "nullable": true, "metadata": {}},
+            {"name": "p_ntz", "type": "timestamp_ntz", "nullable": true, "metadata": {}},
+            {"name": "p_int", "type": "integer", "nullable": true, "metadata": {}},
+            {"name": "value", "type": "integer", "nullable": true, "metadata": {}},
+        ],
+    })
+    .to_string();
+    let protocol = r#"{"protocol":{"minReaderVersion":3,"minWriterVersion":7,"readerFeatures":["timestampNtz"],"writerFeatures":["timestampNtz"]}}"#;
+    let metadata = serde_json::json!({
+        "metaData": {
+            "id": "00000000-0000-0000-0000-000000000000",
+            "format": {"provider": "parquet", "options": {}},
+            "schemaString": schema_string,
+            "partitionColumns": ["p_ts", "p_ntz", "p_int"],
+            "configuration": {"delta.checkpoint.writeStatsAsStruct": "true"},
+            "createdTime": 1700000000000_i64,
+        },
+    })
+    .to_string();
+    add_commit(
+        &table_root,
+        store.as_ref(),
+        0,
+        format!("{protocol}\n{metadata}"),
+    )
+    .await
+    .unwrap();
+    let add = serde_json::json!({
+        "add": {
+            "path": "part.parquet",
+            "partitionValues": {
+                "p_ts": timestamp,
+                "p_ntz": "2024-01-15 12:30:45.123456",
+                "p_int": "7",
+            },
+            "size": 100,
+            "modificationTime": 1700000000000_i64,
+            "dataChange": true,
+            "stats": "{\"numRecords\":1}",
+        },
+    });
+    add_commit(&table_root, store.as_ref(), 1, add.to_string())
+        .await
+        .unwrap();
+    url
+}
+
+#[rstest]
+#[case::utc(None, "2024-01-15 12:30:45.123456", "2024-01-15T12:30:45.123456Z")]
+#[case::los_angeles_winter(
+    Some("America/Los_Angeles"),
+    "2024-01-15 12:30:45.123456",
+    "2024-01-15T20:30:45.123456Z"
+)]
+#[case::los_angeles_summer(
+    Some("America/Los_Angeles"),
+    "2024-06-15 08:00:00.500500",
+    "2024-06-15T15:00:00.500500Z"
+)]
+#[case::fixed_offset_seconds(
+    Some("+12:45:30"),
+    "2024-01-15 12:30:45.123456",
+    "2024-01-14T23:45:15.123456Z"
+)]
+#[case::explicit_offset_wins(
+    Some("America/Los_Angeles"),
+    "2024-01-15 12:30:45+02:00",
+    "2024-01-15T10:30:45Z"
+)]
+#[case::dst_overlap_uses_earlier_instant(
+    Some("America/Los_Angeles"),
+    "2024-11-03 01:30:00",
+    "2024-11-03T08:30:00Z"
+)]
+#[case::dst_gap_uses_pre_transition_offset(
+    Some("America/Los_Angeles"),
+    "2024-03-10 02:30:00",
+    "2024-03-10T10:30:00Z"
+)]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn timezone_aware_parsed_partition_values_across_json_and_native_checkpoint(
+    #[case] timestamp_timezone: Option<&str>,
+    #[case] raw_timestamp: &str,
+    #[case] expected_timestamp: &str,
+    #[values(false, true)] native_checkpoint: bool,
+) {
+    let temp_dir = tempfile::tempdir().unwrap();
+    let url = write_timezone_partition_table(temp_dir.path(), raw_timestamp).await;
+    let engine = create_default_engine_mt_executor(&url).unwrap();
+    if native_checkpoint {
+        let snapshot = Snapshot::builder_for(url.clone())
+            .build(engine.as_ref())
+            .unwrap();
+        snapshot.checkpoint(engine.as_ref(), None).unwrap();
+    }
+
+    let snapshot = Snapshot::builder_for(url).build(engine.as_ref()).unwrap();
+    let partition_values = timestamp_timezone
+        .map_or_else(PartitionValuesOptions::with_struct, |timezone| {
+            PartitionValuesOptions::with_struct().with_timestamp_timezone(timezone)
+        });
+    let scan = snapshot
+        .scan_builder()
+        .with_partition_values(partition_values)
+        .build()
+        .unwrap();
+    let mut rows = 0;
+    for metadata in scan.scan_metadata(engine.as_ref()).unwrap() {
+        let (data, selection) = metadata.unwrap().scan_files.into_parts();
+        let batch: RecordBatch = ArrowEngineData::try_from_engine_data(data).unwrap().into();
+        let batch = filter_record_batch(&batch, &BooleanArray::from(selection)).unwrap();
+        if batch.num_rows() == 0 {
+            continue;
+        }
+        let parsed = get_column!(batch, "partitionValues_parsed", StructArray);
+        let timestamp = parsed
+            .column_by_name("p_ts")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<TimestampMicrosecondArray>()
+            .unwrap();
+        let timestamp_ntz = parsed
+            .column_by_name("p_ntz")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<TimestampMicrosecondArray>()
+            .unwrap();
+        let integer = parsed
+            .column_by_name("p_int")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        assert_eq!(
+            timestamp.value(0),
+            chrono::DateTime::parse_from_rfc3339(expected_timestamp)
+                .unwrap()
+                .timestamp_micros()
+        );
+        assert_eq!(
+            timestamp_ntz.value(0),
+            chrono::DateTime::parse_from_rfc3339("2024-01-15T12:30:45.123456Z")
+                .unwrap()
+                .timestamp_micros(),
+            "reader timezone must not affect TIMESTAMP_NTZ"
+        );
+        assert_eq!(integer.value(0), 7);
+        rows += batch.num_rows();
+    }
+    assert_eq!(rows, 1);
+}
+
+#[rstest]
+#[case::invalid_zone("2024-01-15 12:30:45", "Not/AZone", "Not/AZone")]
+#[case::invalid_timestamp("not a timestamp", "America/Los_Angeles", "not a timestamp")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn timezone_aware_parsed_partition_values_report_invalid_input(
+    #[case] raw_timestamp: &str,
+    #[case] timestamp_timezone: &str,
+    #[case] expected_error: &str,
+) {
+    let temp_dir = tempfile::tempdir().unwrap();
+    let url = write_timezone_partition_table(temp_dir.path(), raw_timestamp).await;
+    let engine = create_default_engine_mt_executor(&url).unwrap();
+    let snapshot = Snapshot::builder_for(url).build(engine.as_ref()).unwrap();
+    let scan = snapshot
+        .scan_builder()
+        .with_partition_values(
+            PartitionValuesOptions::with_struct().with_timestamp_timezone(timestamp_timezone),
+        )
+        .build()
+        .unwrap();
+    let result = scan
+        .scan_metadata(engine.as_ref())
+        .unwrap()
+        .next()
+        .expect("one metadata result");
+    let Err(error) = result else {
+        panic!("invalid timestamp input must fail map parsing");
+    };
+    assert!(error.to_string().contains(expected_error), "{error}");
+}