feat!: support map-to-struct options in engine adapters

Author: DrakeLin

datafusion-executor/src/expression.rs

@@ -253,7 +253,17 @@ uintptr_t convert_engine_to_kernel_expression_item(
       assert(m2s->child_expr.len == 1);
       uintptr_t child = convert_engine_to_kernel_expression_item(
           state, m2s->child_expr.list[0]);
-      return visit_expression_map_to_struct(state, child);
+      struct OptionalValueKernelStringSlice timestamp_timezone = {
+        .tag = NoneKernelStringSlice
+      };
+      if (m2s->timestamp_timezone != NULL) {
+        timestamp_timezone.tag = SomeKernelStringSlice;
+        timestamp_timezone.some = (struct KernelStringSlice) {
+          .ptr = m2s->timestamp_timezone,
+          .len = strlen(m2s->timestamp_timezone)
+        };
+      }
+      return visit_expression_map_to_struct(state, child, timestamp_timezone);
     }
     case StructPatch:
     case FieldPatch:

ffi/examples/visit-expression/engine_to_kernel_expression.h

@@ -131,6 +131,7 @@ struct Column {
 };
 struct MapToStructExpr {
   ExpressionItemList child_expr;
+  char* timestamp_timezone;
 };
 struct BinaryData {
   uint8_t* buf;
@@ -428,9 +429,13 @@ void visit_unknown(void *data, uintptr_t sibling_list_id, struct KernelStringSli
 
 void visit_map_to_struct_expr(void* data,
                               uintptr_t sibling_list_id,
-                              uintptr_t child_list_id) {
+                              uintptr_t child_list_id,
+                              struct OptionalValueKernelStringSlice timestamp_timezone) {
   struct MapToStructExpr* m2s = malloc(sizeof(struct MapToStructExpr));
   m2s->child_expr = get_expr_list(data, child_list_id);
+  m2s->timestamp_timezone = timestamp_timezone.tag == SomeKernelStringSlice
+      ? allocate_string(timestamp_timezone.some)
+      : NULL;
   put_expr_item(data, sibling_list_id, m2s, MapToStruct);
 }
 
@@ -729,6 +734,7 @@ void free_expression_item(ExpressionItem ref) {
     case MapToStruct: {
       struct MapToStructExpr* m2s = ref.ref;
       free_expression_list(m2s->child_expr);
+      free(m2s->timestamp_timezone);
       free(m2s);
       break;
     }

ffi/examples/visit-expression/expression.h

@@ -295,7 +295,11 @@ void print_tree_helper(ExpressionItem ref, int depth) {
     }
     case MapToStruct: {
       struct MapToStructExpr* m2s = ref.ref;
-      printf("MapToStruct\n");
+      if (m2s->timestamp_timezone == NULL) {
+        printf("MapToStruct\n");
+      } else {
+        printf("MapToStruct(timestamp_timezone=%s)\n", m2s->timestamp_timezone);
+      }
       print_expression_item_list(m2s->child_expr, depth + 1);
       break;
     }

ffi/examples/visit-expression/expression_print.h

@@ -16,7 +16,7 @@ use crate::expressions::{
     SharedExpression, SharedOpaqueExpressionOp, SharedOpaquePredicateOp, SharedPredicate,
 };
 use crate::handle::Handle;
-use crate::{kernel_string_slice, KernelStringSlice, SharedSchema};
+use crate::{kernel_string_slice, KernelStringSlice, OptionalValue, SharedSchema};
 
 type VisitLiteralFn<T> = extern "C" fn(data: *mut c_void, sibling_list_id: usize, value: T);
 type VisitUnaryFn = extern "C" fn(data: *mut c_void, sibling_list_id: usize, child_list_id: usize);
@@ -31,6 +31,12 @@ type VisitParseJsonFn = extern "C" fn(
     child_list_id: usize,
     output_schema: Handle<SharedSchema>,
 );
+type VisitMapToStructFn = extern "C" fn(
+    data: *mut c_void,
+    sibling_list_id: usize,
+    child_list_id: usize,
+    timestamp_timezone: OptionalValue<KernelStringSlice>,
+);
 type VisitColumnFn = extern "C" fn(
     data: *mut c_void,
     sibling_list_id: usize,
@@ -174,10 +180,10 @@ pub struct EngineExpressionVisitor {
     /// `child_list_id`. The `output_schema` handle specifies the schema to parse the JSON
     /// into.
     pub visit_parse_json: VisitParseJsonFn,
-    /// Visits a `MapToStruct` expression with default options. Expressions with configured options
-    /// are reported through `visit_unknown` without visiting the child expression. The
-    /// sub-expression is in the one-item list identified by `child_list_id`.
-    pub visit_map_to_struct: VisitUnaryFn,
+    /// Visits a `MapToStruct` expression. The sub-expression is in the one-item list identified by
+    /// `child_list_id`. `timestamp_timezone` carries the configured reader timezone, or is `None`
+    /// when the expression uses the default UTC interpretation.
+    pub visit_map_to_struct: VisitMapToStructFn,
     /// Visits the `LessThan` binary operator belonging to the list identified by
     /// `sibling_list_id`. The operands will be in a _two_ item list identified by
     /// `child_list_id`
@@ -696,13 +702,22 @@ fn visit_expression_impl(
                 schema_handle
             );
         }
-        Expression::MapToStruct(map_to_struct) if !can_visit_map_to_struct(map_to_struct) => {
-            visit_unknown(visitor, sibling_list_id, "configured_map_to_struct")
-        }
-        Expression::MapToStruct(MapToStructExpression { map_expr, .. }) => {
+        Expression::MapToStruct(MapToStructExpression { map_expr, options }) => {
             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);
+            let timestamp_timezone = match options.timestamp_timezone() {
+                Some(timestamp_timezone) => {
+                    OptionalValue::Some(kernel_string_slice!(timestamp_timezone))
+                }
+                None => OptionalValue::None,
+            };
+            call!(
+                visitor,
+                visit_map_to_struct,
+                sibling_list_id,
+                child_list_id,
+                timestamp_timezone
+            );
         }
         // TODO(#2975): Add a dedicated visitor callback for cast expressions.
         Expression::Cast(cast) => visit_unknown(
@@ -714,10 +729,6 @@ fn visit_expression_impl(
     }
 }
 
-fn can_visit_map_to_struct(map_to_struct: &MapToStructExpression) -> bool {
-    map_to_struct.options.is_default()
-}
-
 fn visit_predicate_impl(
     visitor: &mut EngineExpressionVisitor,
     predicate: &Predicate,
@@ -805,6 +816,7 @@ mod tests {
         MapToStruct {
             sibling_list_id: usize,
             child_list_id: usize,
+            timestamp_timezone: Option<String>,
         },
     }
 
@@ -821,6 +833,22 @@ mod tests {
         list_id
     }
 
+    extern "C" fn visit_map_to_struct(
+        data: *mut c_void,
+        sibling_list_id: usize,
+        child_list_id: usize,
+        timestamp_timezone: OptionalValue<KernelStringSlice>,
+    ) {
+        let builder = unsafe { &mut *(data as *mut TestExpressionBuilder) };
+        let timestamp_timezone = Option::from(timestamp_timezone)
+            .map(|timezone| unsafe { String::try_from_slice(&timezone).unwrap() });
+        builder.events.push(LiteralEvent::MapToStruct {
+            sibling_list_id,
+            child_list_id,
+            timestamp_timezone,
+        });
+    }
+
     extern "C" fn visit_literal_interval_year_month(
         data: *mut c_void,
         sibling_list_id: usize,
@@ -875,17 +903,6 @@ mod tests {
         });
     }
 
-    extern "C" fn visit_map_to_struct(
-        data: *mut c_void,
-        sibling_list_id: usize,
-        child_list_id: usize,
-    ) {
-        let builder = unsafe { &mut *(data as *mut TestExpressionBuilder) };
-        builder.events.push(LiteralEvent::MapToStruct {
-            sibling_list_id,
-            child_list_id,
-        });
-    }
     macro_rules! ignore_fn {
         ($fn_name:ident $(, $arg_type:ty)*) => {
             extern "C" fn $fn_name(
@@ -1032,33 +1049,15 @@ mod tests {
         assert_eq!(builder.events, vec![expected]);
     }
 
-    #[test]
-    fn timezone_aware_map_to_struct_visits_unknown() {
-        let expression = Expression::map_to_struct(
-            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: "configured_map_to_struct".to_string(),
-            }]
-        );
-    }
-
-    #[test]
-    fn default_map_to_struct_visits_child_then_map_to_struct() {
-        let expression = Expression::map_to_struct(
-            Expression::column(["partitionValues"]),
-            MapToStructOptions::default(),
-        );
+    #[rstest]
+    #[case::default(None)]
+    #[case::configured(Some("America/Los_Angeles"))]
+    fn map_to_struct_visits_options(#[case] timestamp_timezone: Option<&str>) {
+        let options = timestamp_timezone.map_or_else(MapToStructOptions::default, |timezone| {
+            MapToStructOptions::default().with_timestamp_timezone(timezone)
+        });
+        let expression =
+            Expression::map_to_struct(Expression::column(["partitionValues"]), options);
         let mut builder = TestExpressionBuilder::default();
         let mut visitor = test_visitor(&mut builder);
 
@@ -1075,7 +1074,8 @@ mod tests {
                 LiteralEvent::MapToStruct {
                     sibling_list_id: 0,
                     child_list_id: 1,
-                },
+                    timestamp_timezone: timestamp_timezone.map(str::to_string),
+                }
             ]
         );
     }

ffi/src/expressions/engine_visitor.rs

@@ -20,7 +20,7 @@ use crate::handle::Handle;
 use crate::scan::{EngineExpression, EnginePredicate};
 use crate::{
     AllocateErrorFn, EngineIterator, ExternResult, IntoExternResult, KernelStringSlice,
-    ReferenceSet, TryFromStringSlice,
+    OptionalValue, ReferenceSet, TryFromStringSlice,
 };
 
 pub(crate) enum ExpressionOrPredicate {
@@ -685,17 +685,34 @@ pub extern "C" fn visit_expression_struct(
     wrap_expression(state, Expression::struct_from(exprs))
 }
 
-/// Visit a MapToStruct expression. The `child_expr` is the map expression.
+/// Builds a `MapToStruct` expression from its map child and optional reader timezone.
+///
+/// `timestamp_timezone` is `None` for the default UTC interpretation. A provided string is copied
+/// into the expression before this function returns.
+///
+/// Returns zero when `child_expr` is invalid or the timezone is not valid UTF-8.
+///
+/// # Safety
+///
+/// A provided `timestamp_timezone` slice must have a non-null pointer to a readable buffer of its
+/// declared number of initialized bytes and remain valid for this call.
 #[no_mangle]
-pub extern "C" fn visit_expression_map_to_struct(
+pub unsafe extern "C" fn visit_expression_map_to_struct(
     state: &mut KernelExpressionVisitorState,
     child_expr: usize,
+    timestamp_timezone: OptionalValue<KernelStringSlice>,
 ) -> usize {
+    let options = match Option::from(timestamp_timezone) {
+        Some(timestamp_timezone) => match unsafe { String::try_from_slice(&timestamp_timezone) } {
+            Ok(timestamp_timezone) => {
+                MapToStructOptions::default().with_timestamp_timezone(timestamp_timezone)
+            }
+            Err(_) => return 0,
+        },
+        None => MapToStructOptions::default(),
+    };
     unwrap_kernel_expression(state, child_expr).map_or(0, |expr| {
-        wrap_expression(
-            state,
-            Expression::map_to_struct(expr, MapToStructOptions::default()),
-        )
+        wrap_expression(state, Expression::map_to_struct(expr, options))
     })
 }
 
@@ -871,6 +888,49 @@ mod tests {
 
     use super::*;
 
+    #[rstest]
+    #[case::default(None)]
+    #[case::configured(Some("America/Los_Angeles"))]
+    fn map_to_struct_preserves_options(#[case] timestamp_timezone: Option<&str>) {
+        let mut state = KernelExpressionVisitorState::default();
+        let child = wrap_expression(&mut state, col!("partitionValues"));
+        let ffi_timezone = match timestamp_timezone {
+            Some(timestamp_timezone) => {
+                OptionalValue::Some(crate::kernel_string_slice!(timestamp_timezone))
+            }
+            None => OptionalValue::None,
+        };
+
+        let expression_id =
+            unsafe { visit_expression_map_to_struct(&mut state, child, ffi_timezone) };
+        let expression = unwrap_kernel_expression(&mut state, expression_id).unwrap();
+        let options = timestamp_timezone.map_or_else(MapToStructOptions::default, |timezone| {
+            MapToStructOptions::default().with_timestamp_timezone(timezone)
+        });
+
+        assert_eq!(
+            expression,
+            Expression::map_to_struct(col!("partitionValues"), options)
+        );
+    }
+
+    #[test]
+    fn map_to_struct_rejects_invalid_timezone_utf8() {
+        let mut state = KernelExpressionVisitorState::default();
+        let child = wrap_expression(&mut state, col!("partitionValues"));
+        let invalid_utf8 = [0xff_u8];
+        let timezone = KernelStringSlice {
+            ptr: invalid_utf8.as_ptr().cast(),
+            len: invalid_utf8.len(),
+        };
+
+        let expression_id = unsafe {
+            visit_expression_map_to_struct(&mut state, child, OptionalValue::Some(timezone))
+        };
+
+        assert_eq!(expression_id, 0);
+    }
+
     // ============================================================================
     // NullTypeTag::from_data_type
     // ============================================================================

ffi/src/expressions/kernel_visitor.rs

@@ -160,6 +160,10 @@ pub unsafe extern "C" fn get_testing_kernel_expression() -> Handle<SharedExpress
         Expr::opaque(OpaqueTestOp("foo".to_string()), vec![lit(42), lit(1.111)]),
         Expr::unknown("mystery"),
         Expr::map_to_struct(col!("pv"), MapToStructOptions::default()),
+        Expr::map_to_struct(
+            col!("pv"),
+            MapToStructOptions::default().with_timestamp_timezone("America/Los_Angeles"),
+        ),
         Expr::coalesce([col!("col"), lit(0_i32)]),
         Expr::array([lit(1_i32), lit(2_i32)]),
     ];
@@ -256,6 +260,10 @@ pub unsafe extern "C" fn get_simple_testing_kernel_expression() -> Handle<Shared
         Expr::binary(BinaryExpressionOp::Divide, lit(100), lit(4)),
         Expr::struct_from([lit(1_i32), lit(2_i64), lit(3.0_f64)]),
         Expr::map_to_struct(col!("partitionValues"), MapToStructOptions::default()),
+        Expr::map_to_struct(
+            col!("partitionValues"),
+            MapToStructOptions::default().with_timestamp_timezone("America/Los_Angeles"),
+        ),
     ];
     Arc::new(Expr::struct_from(sub_exprs)).into()
 }

ffi/src/test_ffi.rs

@@ -103,6 +103,8 @@ StructExpression
   Unknown(mystery)
   MapToStruct
     Column(pv)
+  MapToStruct(timestamp_timezone=America/Los_Angeles)
+    Column(pv)
   Coalesce
     Column(col)
     Integer(0)
@@ -216,6 +218,8 @@ StructExpression
     Double(3.000000)
   MapToStruct
     Column(partitionValues)
+  MapToStruct(timestamp_timezone=America/Los_Angeles)
+    Column(partitionValues)
 
 === Expression Round-trip Test ===
 SUCCESS: Round-trip expression matches original!