fix: preserve structured colon path segments

Author: osipovartem

datafusion/functions-nested/src/planner.rs

@@ -93,23 +93,6 @@ 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 }))
@@ -211,104 +194,3 @@ impl ExprPlanner for FieldAccessPlanner {
 fn is_array_agg(func: &Arc<AggregateUDF>) -> 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(_)
-        ));
-    }
-
-    #[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(_)
-            ));
-        }
-    }
-}

datafusion/sql/src/expr/mod.rs

@@ -24,8 +24,8 @@ use datafusion_expr::planner::{
 use sqlparser::ast::{
     AccessExpr, BinaryOperator, CastFormat, CastKind, CeilFloorKind,
     DataType as SQLDataType, DateTimeField, DictionaryField, Expr as SQLExpr,
-    ExprWithAlias as SQLExprWithAlias, JsonPath, MapEntry, Spanned, StructField,
-    Subscript, TrimWhereField, TypedString, Value, ValueWithSpan,
+    ExprWithAlias as SQLExprWithAlias, JsonPath, JsonPathElem, MapEntry, Spanned,
+    StructField, Subscript, TrimWhereField, TypedString, Value, ValueWithSpan,
 };
 use sqlparser::ast::{Query, Visit, Visitor};
 
@@ -839,6 +839,51 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         planner_context: &mut PlannerContext,
         value: Box<SQLExpr>,
         path: &JsonPath,
+    ) -> Result<Expr> {
+        let value = self.sql_to_expr(*value, schema, planner_context)?;
+
+        if !path.path.is_empty()
+            && is_struct_like(&value.get_type(schema)?)
+            && path
+                .path
+                .iter()
+                .all(|element| struct_field_name_from_json_path_elem(element).is_some())
+        {
+            let original_value = value.clone();
+            let mut planned = value;
+
+            for element in &path.path {
+                let field_name = struct_field_name_from_json_path_elem(element)
+                    .expect("all path elements were checked above");
+                let field_access = RawFieldAccessExpr {
+                    expr: planned,
+                    field_access: GetFieldAccess::NamedStructField {
+                        name: ScalarValue::from(field_name),
+                    },
+                };
+                match self.try_plan_field_access(field_access, schema)? {
+                    PlannerResult::Planned(expr) => planned = expr,
+                    PlannerResult::Original(_) => {
+                        return self.plan_serialized_json_access(
+                            original_value,
+                            path,
+                            schema,
+                        );
+                    }
+                }
+            }
+
+            return Ok(planned);
+        }
+
+        self.plan_serialized_json_access(value, path, schema)
+    }
+
+    fn plan_serialized_json_access(
+        &self,
+        value: Expr,
+        path: &JsonPath,
+        schema: &DFSchema,
     ) -> Result<Expr> {
         let json_path = path.to_string();
         let json_path = if let Some(json_path) = json_path.strip_prefix(":") {
@@ -849,13 +894,29 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         };
         self.build_logical_expr(
             BinaryOperator::Custom(":".to_owned()),
-            self.sql_to_expr(*value, schema, planner_context)?,
+            value,
             // pass json path as a string literal, let the impl parse it when needed.
             Expr::Literal(ScalarValue::Utf8(Some(json_path)), None),
             schema,
         )
     }
 
+    fn try_plan_field_access(
+        &self,
+        mut field_access_expr: RawFieldAccessExpr,
+        schema: &DFSchema,
+    ) -> Result<PlannerResult<RawFieldAccessExpr>> {
+        for planner in self.context_provider.get_expr_planners() {
+            match planner.plan_field_access(field_access_expr, schema)? {
+                PlannerResult::Planned(expr) => {
+                    return Ok(PlannerResult::Planned(expr));
+                }
+                PlannerResult::Original(expr) => field_access_expr = expr,
+            }
+        }
+        Ok(PlannerResult::Original(field_access_expr))
+    }
+
     /// Parses a struct(..) expression and plans it creation
     fn parse_struct(
         &self,
@@ -1403,22 +1464,47 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             .into_iter()
             .flatten()
             .try_fold(root, |expr, field_access| {
-                let mut field_access_expr = RawFieldAccessExpr { expr, field_access };
-                for planner in self.context_provider.get_expr_planners() {
-                    match planner.plan_field_access(field_access_expr, schema)? {
-                        PlannerResult::Planned(expr) => return Ok(expr),
-                        PlannerResult::Original(expr) => {
-                            field_access_expr = expr;
-                        }
-                    }
+                let field_access_expr = RawFieldAccessExpr { expr, field_access };
+                match self.try_plan_field_access(field_access_expr, schema)? {
+                    PlannerResult::Planned(expr) => Ok(expr),
+                    PlannerResult::Original(field_access_expr) => not_impl_err!(
+                        "GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
+                    ),
                 }
-                not_impl_err!(
-                    "GetFieldAccess not supported by ExprPlanner: {field_access_expr:?}"
-                )
             })
     }
 }
 
+fn is_struct_like(data_type: &DataType) -> bool {
+    matches!(data_type, DataType::Struct(_))
+        || matches!(
+            data_type,
+            DataType::Dictionary(_, value_type)
+                if matches!(value_type.as_ref(), DataType::Struct(_))
+        )
+}
+
+fn struct_field_name_from_json_path_elem(element: &JsonPathElem) -> Option<&str> {
+    match element {
+        JsonPathElem::Dot { key, .. } => Some(key),
+        JsonPathElem::Bracket {
+            key:
+                SQLExpr::Value(ValueWithSpan {
+                    value: Value::SingleQuotedString(key) | Value::DoubleQuotedString(key),
+                    span: _,
+                }),
+        }
+        | JsonPathElem::ColonBracket {
+            key:
+                SQLExpr::Value(ValueWithSpan {
+                    value: Value::SingleQuotedString(key) | Value::DoubleQuotedString(key),
+                    span: _,
+                }),
+        } => Some(key),
+        JsonPathElem::Bracket { .. } | JsonPathElem::ColonBracket { .. } => None,
+    }
+}
+
 /// Builds a CASE expression that handles NULL semantics for `x <op> ANY(arr)`:
 ///
 /// ```text

datafusion/sqllogictest/test_files/dictionary_struct.slt

@@ -109,6 +109,16 @@ Carol
 Alice
 Bob
 
+# Snowflake-style colon access preserves dictionary encoding and values.
+query TT
+SELECT dict_struct:name, arrow_typeof(dict_struct:name) FROM dict_struct_table;
+----
+Alice Dictionary(UInt32, Utf8)
+Bob Dictionary(UInt32, Utf8)
+Carol Dictionary(UInt32, Utf8)
+Alice Dictionary(UInt32, Utf8)
+Bob Dictionary(UInt32, Utf8)
+
 # A nullable parent makes its non-nullable child nullable, preserving the dictionary type.
 query TTB
 SELECT ds['name'], arrow_typeof(ds['name']), arrow_field(ds['name'])['nullable'] FROM dict_struct_nullable;

datafusion/sqllogictest/test_files/struct.slt

@@ -727,6 +727,22 @@ select get_field(s, 'inner', 'val') from nested_test;
 ----
 100
 
+# Snowflake-style colon access uses the structured JSON path before it is serialized.
+query III
+select
+  s:inner.val as unquoted,
+  s:"inner"."val" as quoted,
+  s:['inner']['val'] as bracketed
+from nested_test;
+----
+100 100 100
+
+# Quoted dots are part of the field name, not path separators.
+query I
+select named_struct('inner.val', 200):"inner.val";
+----
+200
+
 statement ok
 drop table nested_test;