diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index e6fd3230326bc..1a2d323d474d7 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -1043,6 +1043,18 @@ impl Cast { } } +/// Temporary provenance marker for `TRY_CAST` parsed from Snowflake SQL. +pub const EXPLICIT_SQL_TRY_CAST_METADATA_KEY: &str = "datafusion.sql.explicit_try_cast"; + +/// SQL provenance is not a target-field constraint for cast planning. +pub fn cast_metadata_is_type_only( + metadata: &std::collections::HashMap, +) -> bool { + metadata.is_empty() + || (metadata.len() == 1 + && metadata.contains_key(EXPLICIT_SQL_TRY_CAST_METADATA_KEY)) +} + /// TryCast Expression #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] pub struct TryCast { @@ -3634,8 +3646,16 @@ impl Display for Expr { write!(f, "CAST({expr} AS {formatted})") } Expr::TryCast(TryCast { expr, field }) => { - let formatted = - format_type_and_metadata(field.data_type(), Some(field.metadata())); + let formatted = if field + .metadata() + .contains_key(EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + { + let mut metadata = field.metadata().clone(); + metadata.remove(EXPLICIT_SQL_TRY_CAST_METADATA_KEY); + format_type_and_metadata(field.data_type(), Some(&metadata)) + } else { + format_type_and_metadata(field.data_type(), Some(field.metadata())) + }; write!(f, "TRY_CAST({expr} AS {formatted})") } Expr::Not(expr) => write!(f, "NOT {expr}"), diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index a0ec271f55024..4b5d279a87f43 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -18,9 +18,10 @@ use super::{Between, Expr, Like, predicate_bounds}; use crate::ValueOrLambda; use crate::expr::{ - AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Case, Cast, InList, - InSubquery, Lambda, Placeholder, ScalarFunction, TryCast, Unnest, WindowFunction, - WindowFunctionParams, + AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Case, Cast, + EXPLICIT_SQL_TRY_CAST_METADATA_KEY, InList, InSubquery, Lambda, Placeholder, + ScalarFunction, TryCast, Unnest, WindowFunction, WindowFunctionParams, + cast_metadata_is_type_only, }; use crate::expr::{FieldMetadata, LambdaVariable}; use crate::higher_order_function::HigherOrderReturnFieldArgs; @@ -95,9 +96,9 @@ fn cast_output_field( // Check if this is a "type-only" cast (target_field == DataType::X.into_nullable_field()) let is_type_only = target_field.name().is_empty() && target_field.is_nullable() - && target_field.metadata().is_empty(); + && cast_metadata_is_type_only(target_field.metadata()); - let metadata = if is_type_only { + let mut metadata = if is_type_only { // Type-only cast: propagate source metadata, stripping extension type keys let mut meta = source_field.metadata().clone(); meta.remove(EXTENSION_TYPE_NAME_KEY); @@ -107,6 +108,7 @@ fn cast_output_field( // Explicit target field: use target metadata exactly target_field.metadata().clone() }; + metadata.remove(EXPLICIT_SQL_TRY_CAST_METADATA_KEY); let mut f = source_field .as_ref() @@ -1088,6 +1090,32 @@ mod tests { use arrow::datatypes::Schema; use datafusion_common::{DFSchema, assert_or_internal_err}; + #[test] + fn explicit_sql_try_cast_marker_is_not_output_metadata() { + let source = + Arc::new(Field::new("source", DataType::Utf8, false).with_metadata( + HashMap::from([("source_key".to_string(), "x".to_string())]), + )); + let target = Arc::new(Field::new("", DataType::Int64, true).with_metadata( + HashMap::from([( + EXPLICIT_SQL_TRY_CAST_METADATA_KEY.to_string(), + "true".to_string(), + )]), + )); + let output = cast_output_field(&source, &target, true); + assert!( + target + .metadata() + .contains_key(EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + ); + assert!( + !output + .metadata() + .contains_key(EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + ); + assert_eq!(output.metadata().get("source_key"), Some(&"x".to_string())); + } + #[test] fn scalar_function_preserves_udf_coercion_error_source() { #[derive(Debug)] diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index c054026724fb1..6b6714e6febb1 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -31,6 +31,7 @@ use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::{Result, not_impl_err}; use datafusion_expr::ColumnarValue; +use datafusion_expr::expr::cast_metadata_is_type_only; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast #[derive(Debug, Clone, Eq)] @@ -308,7 +309,7 @@ pub fn try_cast_with_target_field( // when only a DataType is known. let is_type_only = target_field.name().is_empty() && target_field.is_nullable() - && target_field.metadata().is_empty(); + && cast_metadata_is_type_only(target_field.metadata()); // For same-type casts, we can skip creating a TryCastExpr only if: // 1. The target is type-only (no explicit metadata) @@ -334,9 +335,19 @@ pub fn try_cast_with_target_field( if is_type_only { Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone()))) } else { + let target_field = if target_field + .metadata() + .contains_key(datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + { + let mut metadata = target_field.metadata().clone(); + metadata.remove(datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY); + Arc::new(target_field.as_ref().clone().with_metadata(metadata)) + } else { + Arc::clone(target_field) + }; Ok(Arc::new(TryCastExpr::new_with_target_field( expr, - Arc::clone(target_field), + target_field, ))) } } @@ -892,6 +903,39 @@ mod tests { Ok(()) } + #[test] + fn explicit_sql_marker_preserves_same_type_fast_path() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let input = col("a", &schema)?; + let target = Arc::new(Field::new("", DataType::Int32, true).with_metadata( + HashMap::from([( + datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY.to_string(), + "true".to_string(), + )]), + )); + let cast = try_cast_with_target_field(Arc::clone(&input), &schema, &target)?; + assert!(Arc::ptr_eq(&input, &cast)); + + let target = Arc::new(Field::new("", DataType::Int64, true).with_metadata( + HashMap::from([ + ( + datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY.to_string(), + "true".to_string(), + ), + ("target_key".to_string(), "x".to_string()), + ]), + )); + let cast = try_cast_with_target_field(input, &schema, &target)?; + let output = cast.return_field(&schema)?; + assert_eq!(output.metadata().get("target_key"), Some(&"x".to_string())); + assert!( + !output + .metadata() + .contains_key(datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + ); + Ok(()) + } + #[test] fn type_only_try_cast_strips_extension_keys() -> Result<()> { // Type-only cast should strip extension keys but preserve other source metadata diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index b1de4e95fd8a2..b2842bf1552d4 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -16,6 +16,7 @@ // under the License. use std::ops::ControlFlow; +use std::sync::Arc; use arrow::datatypes::{DataType, TimeUnit}; use datafusion_expr::planner::{ @@ -29,6 +30,7 @@ use sqlparser::ast::{ }; use sqlparser::ast::{Query, Visit, Visitor}; +use datafusion_common::config::Dialect as SqlDialect; use datafusion_common::{ DFSchema, Diagnostic, Result, ScalarValue, Span, internal_datafusion_err, internal_err, not_impl_err, plan_err, @@ -43,6 +45,7 @@ use datafusion_expr::{ }; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; +use datafusion_expr::expr::EXPLICIT_SQL_TRY_CAST_METADATA_KEY; use datafusion_functions_nested::expr_fn::{ array_has, array_max, array_min, array_position, cardinality, }; @@ -362,7 +365,7 @@ impl SqlToRel<'_, S> { } SQLExpr::Cast { - kind: CastKind::TryCast | CastKind::SafeCast, + kind: kind @ (CastKind::TryCast | CastKind::SafeCast), expr, data_type, format, @@ -372,13 +375,23 @@ impl SqlToRel<'_, S> { return not_impl_err!("CAST with format is not supported: {format}"); } + let mut field = self.convert_data_type_to_field(&data_type)?; + if kind == CastKind::TryCast + && self.context_provider.options().sql_parser.dialect + == SqlDialect::Snowflake + { + let mut metadata = field.metadata().clone(); + metadata + .insert(EXPLICIT_SQL_TRY_CAST_METADATA_KEY.into(), "true".into()); + field = Arc::new(field.as_ref().clone().with_metadata(metadata)); + } Ok(Expr::TryCast(TryCast::new_from_field( Box::new(self.sql_expr_to_logical_expr( *expr, schema, planner_context, )?), - self.convert_data_type_to_field(&data_type)?, + field, ))) } @@ -1572,6 +1585,59 @@ mod tests { test_stack_overflow!(test_stack_overflow_2048, 2048); test_stack_overflow!(test_stack_overflow_4096, 4096); test_stack_overflow!(test_stack_overflow_8192, 8192); + + #[test] + fn snowflake_sql_try_cast_marks_expression_only() { + let dialect = GenericDialect {}; + let mut parser = Parser::new(&dialect) + .try_with_sql("TRY_CAST(column1 AS INTEGER)") + .unwrap(); + let sql_expr = parser.parse_expr().unwrap(); + let schema = DFSchema::try_from_qualified_schema( + "table1", + &Schema::new(vec![Field::new("column1", DataType::Utf8, false)]), + ) + .unwrap(); + + let mut context_provider = TestContextProvider::new(); + context_provider.options.sql_parser.dialect = SqlDialect::Snowflake; + let planned = SqlToRel::new(&context_provider) + .sql_expr_to_logical_expr( + sql_expr.clone(), + &schema, + &mut PlannerContext::default(), + ) + .unwrap(); + let Expr::TryCast(cast) = planned else { + panic!("expected TRY_CAST"); + }; + assert_eq!( + cast.field + .metadata() + .get(EXPLICIT_SQL_TRY_CAST_METADATA_KEY), + Some(&"true".to_string()) + ); + assert!( + !Expr::TryCast(cast.clone()) + .to_string() + .contains(EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + ); + + let generic = TestContextProvider::new(); + let planned = SqlToRel::new(&generic) + .sql_expr_to_logical_expr(sql_expr, &schema, &mut PlannerContext::default()) + .unwrap(); + let Expr::TryCast(cast) = planned else { + panic!("expected TRY_CAST"); + }; + assert!( + !cast + .field + .metadata() + .contains_key(EXPLICIT_SQL_TRY_CAST_METADATA_KEY) + ); + } + #[test] fn test_sql_to_expr_with_alias() { let schema = DFSchema::empty();