feat!: add reader timezone to map-to-struct

Author: DrakeLin

datafusion-executor/src/expression.rs

@@ -16,11 +16,13 @@ use datafusion::logical_expr::{
 };
 use delta_kernel::engine::arrow_conversion::TryIntoArrow;
 use delta_kernel::engine::arrow_data::ArrowEngineData;
+use delta_kernel::engine::arrow_expression::evaluate_expression as kernel_expression;
 use delta_kernel::engine::parse_json;
 use delta_kernel::expressions::{
     BinaryExpression, BinaryExpressionOp, ColumnName as KernelColumnName,
     Expression as KernelExpression, ExpressionRef, ExpressionStructPatch, MapToStructExpression,
-    ParseJsonExpression, UnaryExpressionOp, VariadicExpression, VariadicExpressionOp,
+    MapToStructOptions, ParseJsonExpression, UnaryExpressionOp, VariadicExpression,
+    VariadicExpressionOp,
 };
 use delta_kernel::schema::{
     DataType as KernelDataType, PrimitiveType, SchemaRef as KernelSchemaRef, StructField,
@@ -366,6 +368,11 @@ fn map_to_struct_to_df_expr(
 ) -> DeltaResult<DFExpr> {
     let target = require_struct_output(output_type, "MapToStruct")?;
     let map = to_df_expr(&map_to_struct.map_expr, input_schema, None)?;
+    if let Some(timestamp_timezone) = map_to_struct.options.timestamp_timezone() {
+        let udf =
+            KernelMapToStructUdf::try_new(Arc::new(target.clone()), timestamp_timezone.to_owned())?;
+        return Ok(ScalarUDF::new_from_impl(udf).call(vec![map]));
+    }
 
     let mut args = Vec::with_capacity(target.num_fields() * 2);
     for field in target.fields() {
@@ -394,6 +401,70 @@ fn map_to_struct_to_df_expr(
     Ok(struct_null_when_not(map.is_not_null(), named_struct(args)))
 }
 
+/// A DataFusion scalar UDF that delegates timezone-aware partition parsing to kernel's Arrow
+/// evaluator.
+#[derive(Debug, PartialEq, Eq)]
+struct KernelMapToStructUdf {
+    output_schema: KernelSchemaRef,
+    timestamp_timezone: String,
+    return_type: ArrowDataType,
+    signature: Signature,
+}
+
+impl std::hash::Hash for KernelMapToStructUdf {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        for field in self.output_schema.fields() {
+            field.name().hash(state);
+            field.data_type().to_string().hash(state);
+        }
+        self.timestamp_timezone.hash(state);
+    }
+}
+
+impl KernelMapToStructUdf {
+    fn try_new(output_schema: KernelSchemaRef, timestamp_timezone: String) -> DeltaResult<Self> {
+        let arrow_schema: ArrowSchema = output_schema
+            .as_ref()
+            .try_into_arrow()
+            .map_err(Error::generic_err)?;
+        Ok(Self {
+            return_type: ArrowDataType::Struct(arrow_schema.fields().clone()),
+            output_schema,
+            timestamp_timezone,
+            signature: Signature::any(1, Volatility::Immutable),
+        })
+    }
+}
+
+impl ScalarUDFImpl for KernelMapToStructUdf {
+    fn name(&self) -> &str {
+        "kernel_map_to_struct_with_timestamp_timezone"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[ArrowDataType]) -> Result<ArrowDataType, DataFusionError> {
+        Ok(self.return_type.clone())
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue, DataFusionError> {
+        let num_rows = args.number_rows;
+        let [map] = take_function_args(self.name(), args.args)?;
+        let batch = RecordBatch::try_from_iter([("map", map.into_array(num_rows)?)])?;
+        let expression = KernelExpression::map_to_struct_with_options(
+            KernelExpression::column(["map"]),
+            MapToStructOptions::default().with_timestamp_timezone(self.timestamp_timezone.clone()),
+        );
+        let output_type = KernelDataType::from(self.output_schema.as_ref().clone());
+        let result =
+            kernel_expression::evaluate_expression(&expression, &batch, Some(&output_type))
+                .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        Ok(ColumnarValue::Array(result))
+    }
+}
+
 /// Lowers a `ParseJson` (parse a JSON-string column into a struct) to a call of the
 /// [`ParseJsonUdf`] scalar UDF, which delegates to kernel's own JSON parser. Unlike the
 /// struct-shaped arms, `ParseJson` is self-typed -- it carries its target `output_schema` -- so it
@@ -489,13 +560,17 @@ impl ScalarUDFImpl for ParseJsonUdf {
 
 #[cfg(test)]
 mod tests {
-    use datafusion::arrow::array::{Array, AsArray, StringArray};
+    use datafusion::arrow::array::{
+        Array, AsArray, MapBuilder, StringArray, StringBuilder, TimestampMicrosecondArray,
+    };
+    use datafusion::arrow::datatypes::Field as ArrowField;
     use datafusion::assert_batches_eq;
     use datafusion::common::DFSchema;
     use datafusion::physical_expr::create_physical_expr;
     use datafusion::physical_expr::execution_props::ExecutionProps;
     use delta_kernel::expressions::{
         col, lit, Expression as KernelExpr, ExpressionStructPatch, ExpressionStructPatchBuilder,
+        MapToStructOptions,
     };
     use delta_kernel::schema::{ArrayType, DataType, MapType, StructField, StructType};
     use rstest::rstest;
@@ -1041,6 +1116,55 @@ mod tests {
         assert!(err.contains(expected_message), "{err}");
     }
 
+    #[test]
+    fn map_to_struct_with_timezone_lowers_to_kernel_udf() {
+        let target =
+            StructType::try_new([StructField::nullable("ts", DataType::TIMESTAMP)]).unwrap();
+        let map = KernelExpr::map_to_struct_with_options(
+            col!("pv"),
+            MapToStructOptions::default().with_timestamp_timezone("America/Los_Angeles"),
+        );
+        assert_eq!(
+            to_df_expr(&map, &pv_map_schema(), Some(&DataType::from(target)))
+                .unwrap()
+                .to_string(),
+            "kernel_map_to_struct_with_timestamp_timezone(pv)"
+        );
+    }
+
+    #[test]
+    fn map_to_struct_with_timezone_executes_in_reader_timezone() {
+        let mut maps = MapBuilder::new(None, StringBuilder::new(), StringBuilder::new());
+        maps.keys().append_value("ts");
+        maps.values().append_value("2024-06-15 09:30:00");
+        maps.append(true).unwrap();
+        let map = Arc::new(maps.finish()) as ArrayRef;
+        let arrow_schema =
+            ArrowSchema::new(vec![ArrowField::new("pv", map.data_type().clone(), true)]);
+        let batch = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![map]).unwrap();
+        let target =
+            StructType::try_new([StructField::nullable("ts", DataType::TIMESTAMP)]).unwrap();
+        let logical = to_df_expr(
+            &KernelExpr::map_to_struct_with_options(
+                col!("pv"),
+                MapToStructOptions::default().with_timestamp_timezone("America/Los_Angeles"),
+            ),
+            &pv_map_schema(),
+            Some(&DataType::from(target)),
+        )
+        .unwrap();
+        let df_schema = DFSchema::try_from(arrow_schema).unwrap();
+        let physical = create_physical_expr(&logical, &df_schema, &ExecutionProps::new()).unwrap();
+        let result = physical.evaluate(&batch).unwrap().into_array(1).unwrap();
+        let timestamps = result
+            .as_struct()
+            .column(0)
+            .as_any()
+            .downcast_ref::<TimestampMicrosecondArray>()
+            .unwrap();
+        assert_eq!(timestamps.value(0), 1_718_469_000_000_000);
+    }
+
     // === ParseJson Shared Helpers ===
 
     /// Input schema for JSON tests: `{ j: string }`.

ffi/src/expressions/engine_visitor.rs

@@ -679,7 +679,16 @@ fn visit_expression_impl(
                 schema_handle
             );
         }
-        Expression::MapToStruct(MapToStructExpression { map_expr }) => {
+        Expression::MapToStruct(map_to_struct)
+            if map_to_struct.options.timestamp_timezone().is_some() =>
+        {
+            visit_unknown(
+                visitor,
+                sibling_list_id,
+                "map_to_struct_with_timestamp_timezone",
+            )
+        }
+        Expression::MapToStruct(MapToStructExpression { map_expr, .. }) => {
             let child_list_id = call!(visitor, make_field_list, 1);
             visit_expression_impl(visitor, map_expr, child_list_id);
             call!(visitor, visit_map_to_struct, sibling_list_id, child_list_id);
@@ -754,15 +763,26 @@ fn visit_predicate_internal(predicate: &Predicate, visitor: &mut EngineExpressio
 
 #[cfg(test)]
 mod tests {
-    use delta_kernel::expressions::{Expression, Scalar};
+    use delta_kernel::expressions::{Expression, MapToStructOptions, Scalar};
     use rstest::rstest;
 
     use super::*;
+    use crate::TryFromStringSlice;
 
     #[derive(Debug, PartialEq, Eq)]
     enum LiteralEvent {
-        IntervalYearMonth { sibling_list_id: usize, value: i32 },
-        IntervalDayTime { sibling_list_id: usize, value: i64 },
+        IntervalYearMonth {
+            sibling_list_id: usize,
+            value: i32,
+        },
+        IntervalDayTime {
+            sibling_list_id: usize,
+            value: i64,
+        },
+        Unknown {
+            sibling_list_id: usize,
+            name: String,
+        },
     }
 
     #[derive(Default)]
@@ -802,6 +822,19 @@ mod tests {
         });
     }
 
+    extern "C" fn visit_unknown_name(
+        data: *mut c_void,
+        sibling_list_id: usize,
+        name: KernelStringSlice,
+    ) {
+        let builder = unsafe { &mut *(data as *mut TestExpressionBuilder) };
+        let name = unsafe { String::try_from_slice(&name) }.unwrap();
+        builder.events.push(LiteralEvent::Unknown {
+            sibling_list_id,
+            name,
+        });
+    }
+
     macro_rules! ignore_fn {
         ($fn_name:ident $(, $arg_type:ty)*) => {
             extern "C" fn $fn_name(
@@ -881,7 +914,7 @@ mod tests {
             visit_field_patch: ignore_field_patch,
             visit_opaque_expr: ignore_opaque_expr,
             visit_opaque_pred: ignore_opaque_pred,
-            visit_unknown: ignore_column,
+            visit_unknown: visit_unknown_name,
         }
     }
 
@@ -930,4 +963,25 @@ mod tests {
         assert_eq!(top_level_id, 0);
         assert_eq!(builder.events, vec![expected]);
     }
+
+    #[test]
+    fn timezone_aware_map_to_struct_visits_unknown() {
+        let expression = Expression::map_to_struct_with_options(
+            Expression::column(["partitionValues"]),
+            MapToStructOptions::default().with_timestamp_timezone("America/Los_Angeles"),
+        );
+        let mut builder = TestExpressionBuilder::default();
+        let mut visitor = test_visitor(&mut builder);
+
+        let top_level_id = visit_expression_internal(&expression, &mut visitor);
+
+        assert_eq!(top_level_id, 0);
+        assert_eq!(
+            builder.events,
+            vec![LiteralEvent::Unknown {
+                sibling_list_id: 0,
+                name: "map_to_struct_with_timestamp_timezone".to_string(),
+            }]
+        );
+    }
 }

kernel/proto/expressions.proto

@@ -150,8 +150,15 @@ message ParseJsonExpression {
   delta.kernel.schema.StructType output_schema = 2;
 }
 
+message MapToStructOptions {
+  // Optional IANA timezone or fixed offset for interpreting offset-less TIMESTAMP values. An
+  // explicit offset in a value takes precedence; absence preserves the default UTC behavior.
+  optional string timestamp_timezone = 1;
+}
+
 message MapToStructExpression {
   Expression map_expr = 1;
+  MapToStructOptions options = 2;
 }
 
 // `nullability_predicate` is optional: when set and it evaluates to false/null, the whole