From 220cff6148c9e42e08180d15fb62244b88ebcc85 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Wed, 16 Sep 2026 13:59:04 +0300 Subject: [PATCH 1/2] Plan struct colon access as get_field --- datafusion/functions-nested/src/planner.rs | 86 ++++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/datafusion/functions-nested/src/planner.rs b/datafusion/functions-nested/src/planner.rs index e96fdb7d4baca..39c8617135f1b 100644 --- a/datafusion/functions-nested/src/planner.rs +++ b/datafusion/functions-nested/src/planner.rs @@ -93,6 +93,23 @@ impl ExprPlanner for NestedFunctionPlanner { return Ok(PlannerResult::Planned(array_has_all(right, left))); } } + } else { + #[cfg(feature = "sql")] + let is_colon = + matches!(&op, BinaryOperator::Custom(operator) if operator == ":"); + #[cfg(not(feature = "sql"))] + let is_colon = op == BinaryOperator::Colon; + + if is_colon + && matches!(left.get_type(schema)?, DataType::Struct(_)) + && let Expr::Literal(field_name, _) = &right + && let Some(field_name) = field_name.try_as_str().flatten() + { + return Ok(PlannerResult::Planned(get_field( + left, + field_name.to_owned(), + ))); + } } Ok(PlannerResult::Original(RawBinaryExpr { op, left, right })) @@ -194,3 +211,72 @@ impl ExprPlanner for FieldAccessPlanner { fn is_array_agg(func: &Arc) -> bool { func.name() == "array_agg" } + +#[cfg(all(test, feature = "sql"))] +mod tests { + use super::*; + use arrow::datatypes::{Field, Fields}; + use datafusion_common::{Column, ScalarValue}; + use datafusion_expr::planner::ExprPlanner; + use std::collections::HashMap; + + fn nested_struct_schema() -> DFSchema { + let country = DataType::Struct(Fields::from(vec![Field::new( + "name", + DataType::Utf8, + true, + )])); + let payload = + DataType::Struct(Fields::from(vec![Field::new("country", country, true)])); + DFSchema::from_unqualified_fields( + vec![Field::new("payload", payload, true)].into(), + HashMap::new(), + ) + .unwrap() + } + + fn colon(left: Expr, field_name: &str) -> RawBinaryExpr { + RawBinaryExpr { + op: BinaryOperator::Custom(":".to_owned()), + left, + right: Expr::Literal(ScalarValue::from(field_name), None), + } + } + + fn planned(expr: RawBinaryExpr, schema: &DFSchema) -> Expr { + let PlannerResult::Planned(expr) = + NestedFunctionPlanner.plan_binary_op(expr, schema).unwrap() + else { + panic!("expected colon access to be planned"); + }; + expr + } + + #[test] + fn plans_nested_struct_colon_access_as_get_field() { + let schema = nested_struct_schema(); + let payload = Expr::Column(Column::new_unqualified("payload")); + let country = planned(colon(payload.clone(), "country"), &schema); + let name = planned(colon(country.clone(), "name"), &schema); + + assert_eq!(country, get_field(payload.clone(), "country")); + assert_eq!(name, get_field(get_field(payload, "country"), "name")); + } + + #[test] + fn leaves_non_struct_colon_access_for_other_planners() { + let schema = DFSchema::from_unqualified_fields( + vec![Field::new("text", DataType::Utf8, true)].into(), + HashMap::new(), + ) + .unwrap(); + let original = colon(Expr::Column(Column::new_unqualified("text")), "field"); + + assert!(matches!( + NestedFunctionPlanner + .plan_binary_op(original, &schema) + .unwrap(), + PlannerResult::Original(_) + )); + } +} From 1408bdc34be2941d5c1d9bfa0d66f70b4a952bb5 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Wed, 16 Sep 2026 15:06:36 +0300 Subject: [PATCH 2/2] test: cover unsupported struct colon access --- datafusion/functions-nested/src/planner.rs | 32 ++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/datafusion/functions-nested/src/planner.rs b/datafusion/functions-nested/src/planner.rs index 39c8617135f1b..9590424e57565 100644 --- a/datafusion/functions-nested/src/planner.rs +++ b/datafusion/functions-nested/src/planner.rs @@ -279,4 +279,36 @@ mod tests { PlannerResult::Original(_) )); } + + #[test] + fn leaves_unsupported_struct_access_for_other_planners() { + let schema = nested_struct_schema(); + let payload = Expr::Column(Column::new_unqualified("payload")); + let unsupported = [ + RawBinaryExpr { + op: BinaryOperator::Custom("other".to_owned()), + left: payload.clone(), + right: Expr::Literal(ScalarValue::from("country"), None), + }, + RawBinaryExpr { + op: BinaryOperator::Custom(":".to_owned()), + left: payload.clone(), + right: Expr::Column(Column::new_unqualified("field_name")), + }, + RawBinaryExpr { + op: BinaryOperator::Custom(":".to_owned()), + left: payload, + right: Expr::Literal(ScalarValue::Int64(Some(1)), None), + }, + ]; + + for expression in unsupported { + assert!(matches!( + NestedFunctionPlanner + .plan_binary_op(expression, &schema) + .unwrap(), + PlannerResult::Original(_) + )); + } + } }