diff --git a/datafusion/core/tests/user_defined/expr_planner.rs b/datafusion/core/tests/user_defined/expr_planner.rs index c5e5af731359f..be5fc0775f23e 100644 --- a/datafusion/core/tests/user_defined/expr_planner.rs +++ b/datafusion/core/tests/user_defined/expr_planner.rs @@ -18,6 +18,7 @@ use arrow::array::RecordBatch; use datafusion::common::test_util::batches_to_string; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use datafusion::common::DFSchema; use datafusion::error::Result; @@ -33,6 +34,24 @@ use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr}; #[derive(Debug)] struct MyCustomPlanner; +#[derive(Debug)] +struct DeferColonPlanner { + colon_seen: Arc, +} + +impl ExprPlanner for DeferColonPlanner { + fn plan_binary_op( + &self, + expr: RawBinaryExpr, + _schema: &DFSchema, + ) -> Result> { + if matches!(&expr.op, BinaryOperator::Custom(op) if op == ":") { + self.colon_seen.store(true, Ordering::Relaxed); + } + Ok(PlannerResult::Original(expr)) + } +} + impl ExprPlanner for MyCustomPlanner { fn plan_binary_op( &self, @@ -61,6 +80,13 @@ impl ExprPlanner for MyCustomPlanner { format!("{} ? {}", expr.left, expr.right), )))) } + BinaryOperator::Custom(op) if op == ":" => { + Ok(PlannerResult::Planned(Expr::Alias(Alias::new( + Expr::Literal(ScalarValue::Boolean(Some(true)), None), + None::<&str>, + "custom colon", + )))) + } _ => Ok(PlannerResult::Original(expr)), } } @@ -125,3 +151,54 @@ async fn test_question_filter() { +---+ "); } + +#[tokio::test] +async fn test_custom_struct_colon_operator() { + let config = + SessionConfig::new().set_str("datafusion.sql_parser.dialect", "snowflake"); + let mut ctx = SessionContext::new_with_config(config); + ctx.register_expr_planner(Arc::new(MyCustomPlanner)) + .unwrap(); + + let actual = ctx + .sql("select {'a': 1}:a;") + .await + .unwrap() + .collect() + .await + .unwrap(); + insta::assert_snapshot!(batches_to_string(&actual), @r" + +--------------+ + | custom colon | + +--------------+ + | true | + +--------------+ + "); +} + +#[tokio::test] +async fn test_deferred_struct_colon_operator() { + let config = + SessionConfig::new().set_str("datafusion.sql_parser.dialect", "snowflake"); + let mut ctx = SessionContext::new_with_config(config); + let colon_seen = Arc::new(AtomicBool::new(false)); + ctx.register_expr_planner(Arc::new(DeferColonPlanner { + colon_seen: Arc::clone(&colon_seen), + })) + .unwrap(); + + let actual = ctx + .sql("select {'a': 1}:a;") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert!(colon_seen.load(Ordering::Relaxed)); + assert_eq!(actual.len(), 1); + assert_eq!(actual[0].num_rows(), 1); + assert_eq!( + ScalarValue::try_from_array(actual[0].column(0).as_ref(), 0).unwrap(), + ScalarValue::Int64(Some(1)) + ); +} diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index ea2ed8eea33ce..1372773d35749 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/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}; @@ -313,21 +313,31 @@ impl SqlToRel<'_, S> { right: Expr, schema: &DFSchema, ) -> Result { - // try extension planers - let mut binary_expr = RawBinaryExpr { op, left, right }; + let binary_expr = RawBinaryExpr { op, left, right }; + match self.try_plan_binary_op(binary_expr, schema)? { + PlannerResult::Planned(expr) => Ok(expr), + PlannerResult::Original(RawBinaryExpr { op, left, right }) => { + self.build_binary_expr(&op, left, right) + } + } + } + + fn try_plan_binary_op( + &self, + mut binary_expr: RawBinaryExpr, + schema: &DFSchema, + ) -> Result> { for planner in self.context_provider.get_expr_planners() { match planner.plan_binary_op(binary_expr, schema)? { PlannerResult::Planned(expr) => { - return Ok(expr); + return Ok(PlannerResult::Planned(expr)); } PlannerResult::Original(expr) => { binary_expr = expr; } } } - - let RawBinaryExpr { op, left, right } = binary_expr; - self.build_binary_expr(&op, left, right) + Ok(PlannerResult::Original(binary_expr)) } pub fn sql_to_expr_with_alias( @@ -840,20 +850,72 @@ impl SqlToRel<'_, S> { value: Box, path: &JsonPath, ) -> Result { + let value = self.sql_to_expr(*value, schema, planner_context)?; let json_path = path.to_string(); - let json_path = if let Some(json_path) = json_path.strip_prefix(":") { - // sqlparser's JsonPath display adds an extra `:` at the beginning. - json_path.to_owned() - } else { - json_path + let json_path = json_path.strip_prefix(":").unwrap_or(&json_path); + let binary_expr = RawBinaryExpr { + op: BinaryOperator::Custom(":".to_owned()), + left: value, + right: Expr::Literal(ScalarValue::Utf8(Some(json_path.to_owned())), None), }; - self.build_logical_expr( - BinaryOperator::Custom(":".to_owned()), - self.sql_to_expr(*value, schema, planner_context)?, - // pass json path as a string literal, let the impl parse it when needed. - Expr::Literal(ScalarValue::Utf8(Some(json_path)), None), - schema, - ) + let binary_expr = match self.try_plan_binary_op(binary_expr, schema)? { + PlannerResult::Planned(expr) => return Ok(expr), + PlannerResult::Original(expr) => expr, + }; + + if !path.path.is_empty() + && matches!(&binary_expr.op, BinaryOperator::Custom(op) if op == ":") + && is_struct_like(&binary_expr.left.get_type(schema)?) + && path + .path + .iter() + .all(|element| struct_field_name_from_json_path_elem(element).is_some()) + { + let mut planned = binary_expr.left.clone(); + let mut all_fields_planned = true; + + 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(field_access) => { + planned = field_access.expr; + all_fields_planned = false; + break; + } + } + } + + if all_fields_planned { + return Ok(planned); + } + } + + let RawBinaryExpr { op, left, right } = binary_expr; + self.build_binary_expr(&op, left, right) + } + + fn try_plan_field_access( + &self, + mut field_access_expr: RawFieldAccessExpr, + schema: &DFSchema, + ) -> Result> { + 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 @@ -1403,22 +1465,47 @@ impl 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 ANY(arr)`: /// /// ```text diff --git a/datafusion/sqllogictest/test_files/dictionary_struct.slt b/datafusion/sqllogictest/test_files/dictionary_struct.slt index e2ddc8be1c88c..a4c1312f55e44 100644 --- a/datafusion/sqllogictest/test_files/dictionary_struct.slt +++ b/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; diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index 87bbd11c986a4..66e4f87ed6c55 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/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;