diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index f9c0662e682e8..ba9df50957dd3 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -34,7 +34,6 @@ use crate::{ExprSchemable, Operator, Signature, WindowFrame, WindowUDF}; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cse::{HashNode, NormalizeEq, Normalizeable}; -use datafusion_common::datatype::DataTypeExt; use datafusion_common::metadata::format_type_and_metadata; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion, @@ -984,26 +983,76 @@ pub enum GetFieldAccess { }, } +/// Target of a cast expression. +/// +/// A type-only target inherits metadata from the source expression. An explicit +/// field supplies its own metadata, including an empty map that clears source +/// metadata. +#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] +pub enum CastTarget { + /// Cast to a data type while inheriting source metadata. + DataType(DataType), + /// Cast to an explicit field, including its metadata policy. + Field(FieldRef), +} + +impl CastTarget { + /// Create a target that inherits source metadata. + pub fn type_only(data_type: DataType) -> Self { + Self::DataType(data_type) + } + + /// Create an explicit field target. + pub fn explicit(field: FieldRef) -> Self { + Self::Field(field) + } + + /// Return the target data type. + pub fn data_type(&self) -> &DataType { + match self { + Self::DataType(data_type) => data_type, + Self::Field(field) => field.data_type(), + } + } + + /// Return the explicit target field, or `None` for a type-only target. + pub fn explicit_field(&self) -> Option<&FieldRef> { + match self { + Self::DataType(_) => None, + Self::Field(field) => Some(field), + } + } + + /// Return explicit target metadata, or `None` when metadata is inherited. + pub fn metadata(&self) -> Option<&std::collections::HashMap> { + self.explicit_field().map(|field| field.metadata()) + } +} + /// Cast expression #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] pub struct Cast { /// The expression being cast pub expr: Box, - /// The `DataType` the expression will yield - pub field: FieldRef, + /// The target type and metadata policy. + pub field: CastTarget, } impl Cast { - /// Create a new Cast expression + /// Create a new Cast expression with a type-only target. pub fn new(expr: Box, data_type: DataType) -> Self { Self { expr, - field: data_type.into_nullable_field_ref(), + field: CastTarget::type_only(data_type), } } + /// Create a new Cast expression with explicit target metadata. pub fn new_from_field(expr: Box, field: FieldRef) -> Self { - Self { expr, field } + Self { + expr, + field: CastTarget::explicit(field), + } } } @@ -1012,21 +1061,25 @@ impl Cast { pub struct TryCast { /// The expression being cast pub expr: Box, - /// The `DataType` the expression will yield - pub field: FieldRef, + /// The target type and metadata policy. + pub field: CastTarget, } impl TryCast { - /// Create a new TryCast expression + /// Create a new TryCast expression with a type-only target. pub fn new(expr: Box, data_type: DataType) -> Self { Self { expr, - field: data_type.into_nullable_field_ref(), + field: CastTarget::type_only(data_type), } } + /// Create a new TryCast expression with explicit target metadata. pub fn new_from_field(expr: Box, field: FieldRef) -> Self { - Self { expr, field } + Self { + expr, + field: CastTarget::explicit(field), + } } } @@ -3589,12 +3642,12 @@ impl Display for Expr { } Expr::Cast(Cast { expr, field }) => { let formatted = - format_type_and_metadata(field.data_type(), Some(field.metadata())); + format_type_and_metadata(field.data_type(), field.metadata()); write!(f, "CAST({expr} AS {formatted})") } Expr::TryCast(TryCast { expr, field }) => { let formatted = - format_type_and_metadata(field.data_type(), Some(field.metadata())); + format_type_and_metadata(field.data_type(), field.metadata()); write!(f, "TRY_CAST({expr} AS {formatted})") } Expr::Not(expr) => write!(f, "NOT {expr}"), @@ -4199,10 +4252,10 @@ mod test { #[test] fn format_cast() -> Result<()> { - let expr = Expr::Cast(Cast { - expr: Box::new(Expr::Literal(ScalarValue::Float32(Some(1.23)), None)), - field: DataType::Utf8.into_nullable_field_ref(), - }); + let expr = Expr::Cast(Cast::new( + Box::new(Expr::Literal(ScalarValue::Float32(Some(1.23)), None)), + DataType::Utf8, + )); let expected_canonical = "CAST(Float32(1.23) AS Utf8)"; assert_eq!(expected_canonical, format!("{expr}")); // Note that CAST intentionally has a name that is different from its `Display` diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 36b76f076d26a..fc3e9aa073051 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -18,9 +18,9 @@ use super::{Between, Expr, Like, predicate_bounds}; use crate::ValueOrLambda; use crate::expr::{ - AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Cast, InList, - InSubquery, Lambda, Placeholder, ScalarFunction, TryCast, Unnest, WindowFunction, - WindowFunctionParams, + AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Cast, CastTarget, + InList, InSubquery, Lambda, Placeholder, ScalarFunction, TryCast, Unnest, + WindowFunction, WindowFunctionParams, }; use crate::expr::{FieldMetadata, LambdaVariable}; use crate::higher_order_function::HigherOrderReturnFieldArgs; @@ -71,18 +71,23 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. +/// Derives the output field for a cast expression from the source field, using +/// explicit target metadata when supplied. /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + target: &CastTarget, force_nullable: bool, ) -> Arc { + let metadata = target + .metadata() + .cloned() + .unwrap_or_else(|| source_field.metadata().clone()); let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -623,20 +628,16 @@ impl ExprSchemable for Expr { func.return_field_from_args(args) } // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)), - Expr::Cast(Cast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), false) - }) - } + Expr::Cast(Cast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, false)), Expr::Placeholder(Placeholder { id: _, field: Some(field), }) => Ok(Arc::clone(field).renamed(&schema_name)), - Expr::TryCast(TryCast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), true) - }) - } + Expr::TryCast(TryCast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, true)), Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), @@ -1149,7 +1150,7 @@ mod tests { } #[test] - fn test_expr_metadata() { + fn test_expr_metadata() -> Result<()> { let mut meta = HashMap::new(); meta.insert("bar".to_string(), "buzz".to_string()); let meta = FieldMetadata::from(meta); @@ -1179,6 +1180,45 @@ mod tests { // verify to_field method populates metadata assert_eq!(meta, expr.metadata(&schema).unwrap()); + // An explicit cast target replaces source metadata. A type-only cast + // continues to preserve it. + let target_metadata = HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let target_field = Arc::new( + Field::new("", DataType::FixedSizeBinary(16), true) + .with_metadata(target_metadata.clone()), + ); + let cast = Expr::Cast(Cast::new_from_field( + Box::new(expr.clone()), + Arc::clone(&target_field), + )); + assert_eq!(cast.to_field(&schema)?.1.metadata(), &target_metadata); + + let try_cast = Expr::TryCast(TryCast::new_from_field( + Box::new(expr.clone()), + target_field, + )); + let try_cast_field = try_cast.to_field(&schema)?.1; + assert_eq!(try_cast_field.metadata(), &target_metadata); + assert!(try_cast_field.is_nullable()); + + // An explicitly empty target clears source metadata even when its field + // has the same shape as a synthesized type-only target. + let explicit_empty_target = + Arc::new(Field::new("", DataType::FixedSizeBinary(16), true)); + let cast = Expr::Cast(Cast::new_from_field( + Box::new(expr.clone()), + Arc::clone(&explicit_empty_target), + )); + assert!(cast.to_field(&schema)?.1.metadata().is_empty()); + let try_cast = Expr::TryCast(TryCast::new_from_field( + Box::new(expr.clone()), + explicit_empty_target, + )); + assert!(try_cast.to_field(&schema)?.1.metadata().is_empty()); + // outer ref constructed by `out_ref_col_with_metadata` should be metadata-preserving let outer_ref = out_ref_col_with_metadata( DataType::Int32, @@ -1186,6 +1226,7 @@ mod tests { Column::from_name("foo"), ); assert_eq!(meta, outer_ref.metadata(&schema).unwrap()); + Ok(()) } #[test] diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 75041c701454a..cde5b89ada859 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -112,7 +112,7 @@ pub use datafusion_expr_common::signature::{ }; pub use datafusion_expr_common::type_coercion::binary; pub use expr::{ - Between, BinaryExpr, Case, Cast, Expr, GetFieldAccess, GroupingSet, Like, + Between, BinaryExpr, Case, Cast, CastTarget, Expr, GetFieldAccess, GroupingSet, Like, Sort as SortExpr, TryCast, WindowFunctionDefinition, }; pub use expr_fn::*; diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs index 941fd22ea179f..692ceaafcc170 100644 --- a/datafusion/expr/src/tree_node.rs +++ b/datafusion/expr/src/tree_node.rs @@ -249,10 +249,10 @@ impl TreeNode for Expr { }), Expr::Cast(Cast { expr, field }) => expr .map_elements(f)? - .update_data(|be| Expr::Cast(Cast::new_from_field(be, field))), + .update_data(|expr| Expr::Cast(Cast { expr, field })), Expr::TryCast(TryCast { expr, field }) => expr .map_elements(f)? - .update_data(|be| Expr::TryCast(TryCast::new_from_field(be, field))), + .update_data(|expr| Expr::TryCast(TryCast { expr, field })), Expr::ScalarFunction(ScalarFunction { func, args }) => { args.map_elements(f)?.map_data(|new_args| { Ok(Expr::ScalarFunction(ScalarFunction::new_udf( diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..ebe849d34d460 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -20,9 +20,8 @@ use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::error::ArrowError; use datafusion_common::{ - Result, ScalarValue, arrow_datafusion_err, datatype::DataTypeExt, - exec_datafusion_err, exec_err, internal_err, types::logical_string, - utils::take_function_args, + Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, + internal_err, types::logical_string, utils::take_function_args, }; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; @@ -166,10 +165,10 @@ impl ScalarUDFImpl for ArrowCastFunc { source_arg } else { // Use an actual cast to get the correct type - Expr::Cast(datafusion_expr::Cast { - expr: Box::new(source_arg), - field: target_type.into_nullable_field_ref(), - }) + Expr::Cast(datafusion_expr::Cast::new( + Box::new(source_arg), + target_type, + )) }; // return the newly written argument to DataFusion Ok(ExprSimplifyResult::Simplified(new_expr)) diff --git a/datafusion/functions/src/core/arrow_try_cast.rs b/datafusion/functions/src/core/arrow_try_cast.rs index d27b29ba5736d..107aeaef54732 100644 --- a/datafusion/functions/src/core/arrow_try_cast.rs +++ b/datafusion/functions/src/core/arrow_try_cast.rs @@ -20,8 +20,8 @@ use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::error::ArrowError; use datafusion_common::{ - Result, arrow_datafusion_err, datatype::DataTypeExt, exec_datafusion_err, exec_err, - internal_err, types::logical_string, utils::take_function_args, + Result, arrow_datafusion_err, exec_datafusion_err, exec_err, internal_err, + types::logical_string, utils::take_function_args, }; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; @@ -137,10 +137,10 @@ impl ScalarUDFImpl for ArrowTryCastFunc { let new_expr = if source_type == target_type { source_arg } else { - Expr::TryCast(datafusion_expr::TryCast { - expr: Box::new(source_arg), - field: target_type.into_nullable_field_ref(), - }) + Expr::TryCast(datafusion_expr::TryCast::new( + Box::new(source_arg), + target_type, + )) }; Ok(ExprSimplifyResult::Simplified(new_expr)) } diff --git a/datafusion/functions/src/core/cast_to_type.rs b/datafusion/functions/src/core/cast_to_type.rs index abc7d440e04ba..931bcbc7a2a7e 100644 --- a/datafusion/functions/src/core/cast_to_type.rs +++ b/datafusion/functions/src/core/cast_to_type.rs @@ -130,12 +130,11 @@ impl ScalarUDFImpl for CastToTypeFunc { // the argument's data type is already the correct type source_arg } else { - let nullable = info.nullable(&source_arg)? || target_type == DataType::Null; // Use an actual cast to get the correct type - Expr::Cast(datafusion_expr::Cast { - expr: Box::new(source_arg), - field: Field::new("", target_type, nullable).into(), - }) + Expr::Cast(datafusion_expr::Cast::new( + Box::new(source_arg), + target_type, + )) }; Ok(ExprSimplifyResult::Simplified(new_expr)) } diff --git a/datafusion/functions/src/core/try_cast_to_type.rs b/datafusion/functions/src/core/try_cast_to_type.rs index 4c5af4cc6d228..f4498b1151b89 100644 --- a/datafusion/functions/src/core/try_cast_to_type.rs +++ b/datafusion/functions/src/core/try_cast_to_type.rs @@ -18,9 +18,7 @@ //! [`TryCastToTypeFunc`]: Implementation of the `try_cast_to_type` function use arrow::datatypes::{DataType, Field, FieldRef}; -use datafusion_common::{ - Result, datatype::DataTypeExt, internal_err, utils::take_function_args, -}; +use datafusion_common::{Result, internal_err, utils::take_function_args}; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, @@ -116,10 +114,10 @@ impl ScalarUDFImpl for TryCastToTypeFunc { let new_expr = if source_type == target_type { source_arg } else { - Expr::TryCast(datafusion_expr::TryCast { - expr: Box::new(source_arg), - field: target_type.into_nullable_field_ref(), - }) + Expr::TryCast(datafusion_expr::TryCast::new( + Box::new(source_arg), + target_type, + )) }; Ok(ExprSimplifyResult::Simplified(new_expr)) } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index cb3103d38c52a..220a9fae7b63f 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -54,13 +54,48 @@ fn can_cast_named_struct_types(source: &DataType, target: &DataType) -> bool { validate_data_type_compatibility("", source, target).is_ok() } +/// Target type and metadata policy shared by physical cast expressions. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) enum PhysicalCastTarget { + /// A synthesized target that inherits source field properties. + DataType(FieldRef), + /// An explicit target field, including explicitly empty metadata. + Field(FieldRef), +} + +impl PhysicalCastTarget { + pub(crate) fn type_only(data_type: DataType) -> Self { + Self::DataType(data_type.into_nullable_field_ref()) + } + + pub(crate) fn explicit(field: FieldRef) -> Self { + Self::Field(field) + } + + pub(crate) fn data_type(&self) -> &DataType { + self.field().data_type() + } + + pub(crate) fn field(&self) -> &FieldRef { + match self { + Self::DataType(field) | Self::Field(field) => field, + } + } + + pub(crate) fn explicit_field(&self) -> Option<&FieldRef> { + match self { + Self::DataType(_) => None, + Self::Field(field) => Some(field), + } + } +} + /// CAST expression casts an expression to a specific data type and returns a runtime error on invalid cast #[derive(Debug, Clone, Eq)] pub struct CastExpr { /// The expression to cast pub expr: Arc, - /// Field metadata describing the desired output after casting - target_field: FieldRef, + target: PhysicalCastTarget, /// Cast options cast_options: CastOptions<'static>, } @@ -69,7 +104,7 @@ pub struct CastExpr { impl PartialEq for CastExpr { fn eq(&self, other: &Self) -> bool { self.expr.eq(&other.expr) - && self.target_field.eq(&other.target_field) + && self.target == other.target && self.cast_options.eq(&other.cast_options) } } @@ -77,7 +112,7 @@ impl PartialEq for CastExpr { impl Hash for CastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.target_field.hash(state); + self.target.hash(state); self.cast_options.hash(state); } } @@ -102,9 +137,9 @@ impl CastExpr { cast_type: DataType, cast_options: Option>, ) -> Self { - Self::new_with_target_field( + Self::new_with_target( expr, - cast_type.into_nullable_field_ref(), + PhysicalCastTarget::type_only(cast_type), cast_options, ) } @@ -122,10 +157,22 @@ impl CastExpr { expr: Arc, target_field: FieldRef, cast_options: Option>, + ) -> Self { + Self::new_with_target( + expr, + PhysicalCastTarget::explicit(target_field), + cast_options, + ) + } + + fn new_with_target( + expr: Arc, + target: PhysicalCastTarget, + cast_options: Option>, ) -> Self { Self { expr, - target_field, + target, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), } } @@ -137,12 +184,12 @@ impl CastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - self.target_field.data_type() + self.target.data_type() } - /// Field metadata describing the output column after casting. + /// The field stored in the cast target. pub fn target_field(&self) -> &FieldRef { - &self.target_field + self.target.field() } /// The cast options @@ -151,20 +198,24 @@ impl CastExpr { } fn resolved_target_field(&self, input_schema: &Schema) -> Result { - if is_default_target_field(&self.target_field) { - self.expr.return_field(input_schema).map(|field| { + match self.target.explicit_field() { + Some(field) => Ok(Arc::clone(field)), + None => self.expr.return_field(input_schema).map(|source| { Arc::new( - field + source .as_ref() .clone() .with_data_type(self.cast_type().clone()), ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + }), } } + /// Return this cast with a new input expression, preserving its target. + pub fn with_new_expr(&self, expr: Arc) -> Self { + Self::new_with_target(expr, self.target.clone(), Some(self.cast_options.clone())) + } + /// Check if casting from the specified source type to the target type is a /// widening cast (e.g. from `Int8` to `Int16`). pub fn check_bigger_cast(cast_type: &DataType, src: &DataType) -> bool { @@ -191,12 +242,6 @@ impl CastExpr { } } -fn is_default_target_field(target_field: &FieldRef) -> bool { - target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty() -} - pub(crate) fn is_order_preserving_cast_family( source_type: &DataType, target_type: &DataType, @@ -267,11 +312,7 @@ impl PhysicalExpr for CastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(CastExpr::new_with_target_field( - Arc::clone(&children[0]), - Arc::clone(&self.target_field), - Some(self.cast_options.clone()), - ))) + Ok(Arc::new(self.with_new_expr(Arc::clone(&children[0])))) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -319,6 +360,11 @@ impl PhysicalExpr for CastExpr { protobuf::PhysicalCastNode { expr: Some(Box::new(ctx.encode_child(self.expr())?)), arrow_type: Some(self.cast_type().try_into()?), + target_field: self + .target + .explicit_field() + .map(|field| field.as_ref().try_into()) + .transpose()?, }, ))), })) @@ -358,7 +404,23 @@ impl CastExpr { internal_datafusion_err!("CastExpr is missing required field 'arrow_type'") })?; - Ok(Arc::new(CastExpr::new(expr, arrow_type.try_into()?, None))) + let data_type: DataType = arrow_type.try_into()?; + if let Some(target_field) = cast_expr.target_field.as_ref() { + let field: arrow::datatypes::Field = target_field.try_into()?; + if field.data_type() != &data_type { + return internal_err!( + "CastExpr target_field type {} does not match arrow_type {data_type}", + field.data_type() + ); + } + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + Arc::new(field), + None, + ))) + } else { + Ok(Arc::new(CastExpr::new(expr, data_type, None))) + } } } @@ -372,12 +434,21 @@ pub fn cast_with_options( cast_type: DataType, cast_options: Option>, ) -> Result> { - cast_with_target_field( - expr, - input_schema, - cast_type.into_nullable_field_ref(), - cast_options, - ) + let expr_type = expr.data_type(input_schema)?; + if expr_type == cast_type { + return Ok(Arc::clone(&expr)); + } + + let can_build_cast = if requires_nested_struct_cast(&expr_type, &cast_type) { + can_cast_named_struct_types(&expr_type, &cast_type) + } else { + can_cast_types(&expr_type, &cast_type) + }; + if !can_build_cast { + return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); + } + + Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options))) } /// Return a PhysicalExpression representing `expr` casted to `target_field`, @@ -396,10 +467,6 @@ pub fn cast_with_target_field( ) -> Result> { let expr_type = expr.data_type(input_schema)?; let cast_type = target_field.data_type(); - if expr_type == *cast_type && is_default_target_field(&target_field) { - return Ok(Arc::clone(&expr)); - } - let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { // Allow casts involving structs (including nested inside Lists, Dictionaries, // etc.) that pass name-based compatibility validation. This validation is @@ -1277,7 +1344,11 @@ mod proto_tests { PhysicalExprNode { expr_id: None, expr_type: Some(physical_expr_node::ExprType::Cast(Box::new( - PhysicalCastNode { expr, arrow_type }, + PhysicalCastNode { + expr, + arrow_type, + target_field: None, + }, ))), } } @@ -1308,6 +1379,52 @@ mod proto_tests { assert_eq!(data_type, Int64); } + #[test] + fn proto_roundtrip_preserves_explicit_cast_target() { + let target = Arc::new(Field::new("", Int64, true)); + let schema = Schema::new(vec![Field::new("a", Int32, false)]); + let cast = CastExpr::new_with_target_field( + col("a", &schema).unwrap(), + Arc::clone(&target), + None, + ); + let encoded = cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&StubEncoder::ok())) + .unwrap() + .unwrap(); + let Some(physical_expr_node::ExprType::Cast(encoded)) = encoded.expr_type else { + unreachable!() + }; + assert_eq!( + encoded.target_field, + Some(target.as_ref().try_into().unwrap()) + ); + + let mut node = proto_cast_node( + Some(Box::new(column_node("a"))), + Some(proto_int64_arrow_type()), + ); + let Some(physical_expr_node::ExprType::Cast(proto_cast)) = + node.expr_type.as_mut() + else { + unreachable!() + }; + proto_cast.target_field = encoded.target_field; + let decoded = CastExpr::try_from_proto( + &node, + &PhysicalExprDecodeCtx::new(&Schema::empty(), &StubDecoder::ok()), + ) + .unwrap(); + assert_eq!( + decoded + .downcast_ref::() + .unwrap() + .target + .explicit_field(), + Some(&target) + ); + } + #[test] fn try_to_proto_propagates_child_encode_error() { let cast = proto_cast_fixture(); diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 521b8b87e305c..f2f9285de560a 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -61,7 +61,7 @@ pub use no_op::NoOp; pub use not::{NotExpr, not}; pub(crate) use similar_to_pattern::translate_scalar; pub use similar_to_pattern::{SqlSimilarToPattern, sql_similar_to_regex}; -pub use try_cast::{TryCastExpr, try_cast}; +pub use try_cast::{TryCastExpr, try_cast, try_cast_with_target_field}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..07ad1a326c89a 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -20,6 +20,7 @@ use std::hash::Hash; use std::sync::Arc; use crate::PhysicalExpr; +use crate::expressions::cast::PhysicalCastTarget; use arrow::compute; use arrow::compute::CastOptions; use arrow::datatypes::{DataType, FieldRef, Schema}; @@ -34,28 +35,42 @@ use datafusion_expr::ColumnarValue; pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + target: PhysicalCastTarget, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for TryCastExpr { fn eq(&self, other: &Self) -> bool { - self.expr.eq(&other.expr) && self.cast_type == other.cast_type + self.expr.eq(&other.expr) && self.target == other.target } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.target.hash(state); } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self { + expr, + target: PhysicalCastTarget::type_only(cast_type), + } + } + + /// Create a new `TryCastExpr` with an explicit target field, preserving its + /// name and metadata. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { + expr, + target: PhysicalCastTarget::explicit(target_field), + } } /// The expression to cast @@ -65,19 +80,27 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target.data_type() + } + + /// Return this cast with a new input expression, preserving its target. + pub fn with_new_expr(&self, expr: Arc) -> Self { + Self { + expr, + target: self.target.clone(), + } } } impl fmt::Display for TryCastExpr { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type) + write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type()) } } impl PhysicalExpr for TryCastExpr { fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(self.cast_type.clone()) + Ok(self.cast_type().clone()) } fn nullable(&self, _input_schema: &Schema) -> Result { @@ -90,14 +113,20 @@ impl PhysicalExpr for TryCastExpr { safe: true, format_options: DEFAULT_FORMAT_OPTIONS, }; - value.cast_to(&self.cast_type, Some(&options)) + value.cast_to(self.cast_type(), Some(&options)) } fn return_field(&self, input_schema: &Schema) -> Result { - self.expr - .return_field(input_schema) - .map(|f| f.as_ref().clone().with_data_type(self.cast_type.clone())) - .map(Arc::new) + let field = if let Some(target) = self.target.explicit_field() { + target.as_ref().clone() + } else { + self.expr + .return_field(input_schema)? + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + }; + Ok(Arc::new(field.with_nullable(true))) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +137,13 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( - Arc::clone(&children[0]), - self.cast_type.clone(), - ))) + Ok(Arc::new(self.with_new_expr(Arc::clone(&children[0])))) } fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "TRY_CAST(")?; self.expr.fmt_sql(f)?; - write!(f, " AS {:?})", self.cast_type) + write!(f, " AS {:?})", self.cast_type()) } #[cfg(feature = "proto")] @@ -133,6 +159,11 @@ impl PhysicalExpr for TryCastExpr { protobuf::PhysicalTryCastNode { expr: Some(Box::new(ctx.encode_child(&self.expr)?)), arrow_type: Some(self.cast_type().try_into()?), + target_field: self + .target + .explicit_field() + .map(|field| field.as_ref().try_into()) + .transpose()?, }, ))), })) @@ -165,9 +196,22 @@ impl TryCastExpr { "TryCastExpr", "arrow_type", )?; - let cast_type: DataType = arrow_type.try_into()?; - - Ok(Arc::new(TryCastExpr::new(expr, cast_type))) + let data_type: DataType = arrow_type.try_into()?; + if let Some(target_field) = try_cast.target_field.as_ref() { + let field: arrow::datatypes::Field = target_field.try_into()?; + if field.data_type() != &data_type { + return datafusion_common::internal_err!( + "TryCastExpr target_field type {} does not match arrow_type {data_type}", + field.data_type() + ); + } + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + Arc::new(field), + ))) + } else { + Ok(Arc::new(TryCastExpr::new(expr, data_type))) + } } } @@ -190,6 +234,25 @@ pub fn try_cast( } } +/// Return a physical expression that tries to cast `expr` to an explicit target +/// field, preserving its name and metadata even when the data type already matches. +pub fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: FieldRef, +) -> Result> { + let expr_type = expr.data_type(input_schema)?; + let cast_type = target_field.data_type(); + if can_cast_types(&expr_type, cast_type) { + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + target_field, + ))) + } else { + not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}") + } +} + #[cfg(test)] mod tests { use super::*; @@ -676,7 +739,11 @@ mod proto_tests { PhysicalExprNode { expr_id: None, expr_type: Some(physical_expr_node::ExprType::TryCast(Box::new( - PhysicalTryCastNode { expr, arrow_type }, + PhysicalTryCastNode { + expr, + arrow_type, + target_field: None, + }, ))), } } @@ -707,6 +774,50 @@ mod proto_tests { assert_eq!(data_type, DataType::Int32); } + #[test] + fn proto_roundtrip_preserves_explicit_try_cast_target() { + let target = Arc::new(Field::new("", DataType::Int32, true)); + let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]); + let try_cast = TryCastExpr::new_with_target_field( + col("a", &schema).unwrap(), + Arc::clone(&target), + ); + let encoded = try_cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&StubEncoder::ok())) + .unwrap() + .unwrap(); + let Some(physical_expr_node::ExprType::TryCast(encoded)) = encoded.expr_type + else { + unreachable!() + }; + assert_eq!( + encoded.target_field, + Some(target.as_ref().try_into().unwrap()) + ); + + let mut node = + try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type())); + let Some(physical_expr_node::ExprType::TryCast(proto_try_cast)) = + node.expr_type.as_mut() + else { + unreachable!() + }; + proto_try_cast.target_field = encoded.target_field; + let decoded = TryCastExpr::try_from_proto( + &node, + &PhysicalExprDecodeCtx::new(&Schema::empty(), &StubDecoder::ok()), + ) + .unwrap(); + assert_eq!( + decoded + .downcast_ref::() + .unwrap() + .target + .explicit_field(), + Some(&target) + ); + } + #[test] fn try_to_proto_propagates_child_encode_error() { let try_cast = try_cast_fixture(); diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 9a0bdc33da8e9..cffc4ffd99bae 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -384,14 +384,30 @@ pub fn create_physical_expr( }; Ok(expressions::case(expr, when_then_expr, else_expr)?) } - Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( - create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, - input_schema, - Arc::clone(field), - None, - ), + Expr::Cast(Cast { expr, field }) => { + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + if field.explicit_field().is_some() { + let (_, output_field) = e.to_field(input_dfschema)?; + expressions::cast_with_target_field( + physical_expr, + input_schema, + output_field, + None, + ) + } else { + expressions::cast(physical_expr, input_schema, field.data_type().clone()) + } + } Expr::TryCast(TryCast { expr, field }) => { - if !field.metadata().is_empty() { + if field + .metadata() + .is_some_and(|metadata| !metadata.is_empty()) + { let (_, src_field) = expr.to_field(input_dfschema)?; return plan_err!( "TryCast from {} to {} is not supported", @@ -399,20 +415,30 @@ pub fn create_physical_expr( src_field.data_type(), Some(src_field.metadata()), ), - format_type_and_metadata(field.data_type(), Some(field.metadata())) + format_type_and_metadata(field.data_type(), field.metadata()) ); } - expressions::try_cast( - create_physical_expr( - expr, - input_dfschema, - execution_props, - planning_ctx, - )?, - input_schema, - field.data_type().clone(), - ) + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + if field.explicit_field().is_some() { + let (_, output_field) = e.to_field(input_dfschema)?; + expressions::try_cast_with_target_field( + physical_expr, + input_schema, + output_field, + ) + } else { + expressions::try_cast( + physical_expr, + input_schema, + field.data_type().clone(), + ) + } } Expr::Not(expr) => expressions::not(create_physical_expr( expr, @@ -842,12 +868,17 @@ mod tests { Arc::clone(&target_field), )); + let logical_output = cast_expr.to_field(&DFSchema::try_from(schema.clone())?)?.1; let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); + let output = physical.return_field(&schema)?; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + assert_eq!(output, logical_output); + assert_eq!(cast.target_field(), &output); + assert_eq!(output.name(), "a"); + assert_eq!(output.data_type(), &DataType::Int64); + assert_eq!(output.metadata(), target_field.metadata()); + assert!(!physical.nullable(&schema)?); Ok(()) } @@ -884,10 +915,12 @@ mod tests { let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); + let output = physical.return_field(&schema)?; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + assert_eq!(cast.target_field(), &output); + assert_eq!(output.name(), "a"); + assert_eq!(output.metadata(), target_field.metadata()); + assert!(!physical.nullable(&schema)?); Ok(()) } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index ce5d11425c52d..796ad0ab556fb 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -738,6 +738,7 @@ message CastNode { datafusion_common.ArrowType arrow_type = 2; map metadata = 3; optional bool nullable = 4; + datafusion_common.Field target_field = 5; } message TryCastNode { @@ -745,6 +746,7 @@ message TryCastNode { datafusion_common.ArrowType arrow_type = 2; map metadata = 3; optional bool nullable = 4; + datafusion_common.Field target_field = 5; } message SortExprNode { @@ -1186,11 +1188,13 @@ message PhysicalCaseNode { message PhysicalTryCastNode { PhysicalExprNode expr = 1; datafusion_common.ArrowType arrow_type = 2; + datafusion_common.Field target_field = 3; } message PhysicalCastNode { PhysicalExprNode expr = 1; datafusion_common.ArrowType arrow_type = 2; + datafusion_common.Field target_field = 3; } message PhysicalNegativeNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 0f622de3fa89f..ec64e512085ad 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -2446,6 +2446,9 @@ impl serde::Serialize for CastNode { if self.nullable.is_some() { len += 1; } + if self.target_field.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.CastNode", len)?; if let Some(v) = self.expr.as_ref() { struct_ser.serialize_field("expr", v)?; @@ -2459,6 +2462,9 @@ impl serde::Serialize for CastNode { if let Some(v) = self.nullable.as_ref() { struct_ser.serialize_field("nullable", v)?; } + if let Some(v) = self.target_field.as_ref() { + struct_ser.serialize_field("targetField", v)?; + } struct_ser.end() } } @@ -2474,6 +2480,8 @@ impl<'de> serde::Deserialize<'de> for CastNode { "arrowType", "metadata", "nullable", + "target_field", + "targetField", ]; #[allow(clippy::enum_variant_names)] @@ -2482,6 +2490,7 @@ impl<'de> serde::Deserialize<'de> for CastNode { ArrowType, Metadata, Nullable, + TargetField, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -2507,6 +2516,7 @@ impl<'de> serde::Deserialize<'de> for CastNode { "arrowType" | "arrow_type" => Ok(GeneratedField::ArrowType), "metadata" => Ok(GeneratedField::Metadata), "nullable" => Ok(GeneratedField::Nullable), + "targetField" | "target_field" => Ok(GeneratedField::TargetField), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -2530,6 +2540,7 @@ impl<'de> serde::Deserialize<'de> for CastNode { let mut arrow_type__ = None; let mut metadata__ = None; let mut nullable__ = None; + let mut target_field__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -2558,6 +2569,12 @@ impl<'de> serde::Deserialize<'de> for CastNode { } nullable__ = map_.next_value()?; } + GeneratedField::TargetField => { + if target_field__.is_some() { + return Err(serde::de::Error::duplicate_field("targetField")); + } + target_field__ = map_.next_value()?; + } } } Ok(CastNode { @@ -2565,6 +2582,7 @@ impl<'de> serde::Deserialize<'de> for CastNode { arrow_type: arrow_type__, metadata: metadata__.unwrap_or_default(), nullable: nullable__, + target_field: target_field__, }) } } @@ -18048,6 +18066,9 @@ impl serde::Serialize for PhysicalCastNode { if self.arrow_type.is_some() { len += 1; } + if self.target_field.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalCastNode", len)?; if let Some(v) = self.expr.as_ref() { struct_ser.serialize_field("expr", v)?; @@ -18055,6 +18076,9 @@ impl serde::Serialize for PhysicalCastNode { if let Some(v) = self.arrow_type.as_ref() { struct_ser.serialize_field("arrowType", v)?; } + if let Some(v) = self.target_field.as_ref() { + struct_ser.serialize_field("targetField", v)?; + } struct_ser.end() } } @@ -18068,12 +18092,15 @@ impl<'de> serde::Deserialize<'de> for PhysicalCastNode { "expr", "arrow_type", "arrowType", + "target_field", + "targetField", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Expr, ArrowType, + TargetField, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -18097,6 +18124,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCastNode { match value { "expr" => Ok(GeneratedField::Expr), "arrowType" | "arrow_type" => Ok(GeneratedField::ArrowType), + "targetField" | "target_field" => Ok(GeneratedField::TargetField), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18118,6 +18146,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCastNode { { let mut expr__ = None; let mut arrow_type__ = None; + let mut target_field__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -18132,11 +18161,18 @@ impl<'de> serde::Deserialize<'de> for PhysicalCastNode { } arrow_type__ = map_.next_value()?; } + GeneratedField::TargetField => { + if target_field__.is_some() { + return Err(serde::de::Error::duplicate_field("targetField")); + } + target_field__ = map_.next_value()?; + } } } Ok(PhysicalCastNode { expr: expr__, arrow_type: arrow_type__, + target_field: target_field__, }) } } @@ -21998,6 +22034,9 @@ impl serde::Serialize for PhysicalTryCastNode { if self.arrow_type.is_some() { len += 1; } + if self.target_field.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalTryCastNode", len)?; if let Some(v) = self.expr.as_ref() { struct_ser.serialize_field("expr", v)?; @@ -22005,6 +22044,9 @@ impl serde::Serialize for PhysicalTryCastNode { if let Some(v) = self.arrow_type.as_ref() { struct_ser.serialize_field("arrowType", v)?; } + if let Some(v) = self.target_field.as_ref() { + struct_ser.serialize_field("targetField", v)?; + } struct_ser.end() } } @@ -22018,12 +22060,15 @@ impl<'de> serde::Deserialize<'de> for PhysicalTryCastNode { "expr", "arrow_type", "arrowType", + "target_field", + "targetField", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Expr, ArrowType, + TargetField, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -22047,6 +22092,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalTryCastNode { match value { "expr" => Ok(GeneratedField::Expr), "arrowType" | "arrow_type" => Ok(GeneratedField::ArrowType), + "targetField" | "target_field" => Ok(GeneratedField::TargetField), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -22068,6 +22114,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalTryCastNode { { let mut expr__ = None; let mut arrow_type__ = None; + let mut target_field__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -22082,11 +22129,18 @@ impl<'de> serde::Deserialize<'de> for PhysicalTryCastNode { } arrow_type__ = map_.next_value()?; } + GeneratedField::TargetField => { + if target_field__.is_some() { + return Err(serde::de::Error::duplicate_field("targetField")); + } + target_field__ = map_.next_value()?; + } } } Ok(PhysicalTryCastNode { expr: expr__, arrow_type: arrow_type__, + target_field: target_field__, }) } } @@ -26888,6 +26942,9 @@ impl serde::Serialize for TryCastNode { if self.nullable.is_some() { len += 1; } + if self.target_field.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.TryCastNode", len)?; if let Some(v) = self.expr.as_ref() { struct_ser.serialize_field("expr", v)?; @@ -26901,6 +26958,9 @@ impl serde::Serialize for TryCastNode { if let Some(v) = self.nullable.as_ref() { struct_ser.serialize_field("nullable", v)?; } + if let Some(v) = self.target_field.as_ref() { + struct_ser.serialize_field("targetField", v)?; + } struct_ser.end() } } @@ -26916,6 +26976,8 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { "arrowType", "metadata", "nullable", + "target_field", + "targetField", ]; #[allow(clippy::enum_variant_names)] @@ -26924,6 +26986,7 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { ArrowType, Metadata, Nullable, + TargetField, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -26949,6 +27012,7 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { "arrowType" | "arrow_type" => Ok(GeneratedField::ArrowType), "metadata" => Ok(GeneratedField::Metadata), "nullable" => Ok(GeneratedField::Nullable), + "targetField" | "target_field" => Ok(GeneratedField::TargetField), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -26972,6 +27036,7 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { let mut arrow_type__ = None; let mut metadata__ = None; let mut nullable__ = None; + let mut target_field__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -27000,6 +27065,12 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { } nullable__ = map_.next_value()?; } + GeneratedField::TargetField => { + if target_field__.is_some() { + return Err(serde::de::Error::duplicate_field("targetField")); + } + target_field__ = map_.next_value()?; + } } } Ok(TryCastNode { @@ -27007,6 +27078,7 @@ impl<'de> serde::Deserialize<'de> for TryCastNode { arrow_type: arrow_type__, metadata: metadata__.unwrap_or_default(), nullable: nullable__, + target_field: target_field__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 7527c588954c5..0af1d6ae48de1 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1178,6 +1178,8 @@ pub struct CastNode { >, #[prost(bool, optional, tag = "4")] pub nullable: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct TryCastNode { @@ -1192,6 +1194,8 @@ pub struct TryCastNode { >, #[prost(bool, optional, tag = "4")] pub nullable: ::core::option::Option, + #[prost(message, optional, tag = "5")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SortExprNode { @@ -1842,6 +1846,8 @@ pub struct PhysicalTryCastNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, tag = "2")] pub arrow_type: ::core::option::Option, + #[prost(message, optional, tag = "3")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalCastNode { @@ -1849,6 +1855,8 @@ pub struct PhysicalCastNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, tag = "2")] pub arrow_type: ::core::option::Option, + #[prost(message, optional, tag = "3")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalNegativeNode { diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index d4d0ea7292ffe..197aef5571db0 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -18,7 +18,6 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field}; -use datafusion_common::datatype::DataTypeExt; use datafusion_common::{ Result, ScalarValue, SplitPoint, TableReference, exec_datafusion_err, internal_err, plan_datafusion_err, @@ -34,7 +33,7 @@ use datafusion_expr::expr::{ use datafusion_expr::expr::{Unnest, WildcardOptions}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ - Between, BinaryExpr, Case, Cast, Expr, GroupingSet, + Between, BinaryExpr, Case, Cast, CastTarget, Expr, GroupingSet, GroupingSet::GroupingSets, Like, Operator, TryCast, WindowFrame, expr::{self, InList, WindowFunction}, @@ -46,6 +45,35 @@ use crate::protobuf::{self, CubeNode, GroupingSetNode, PlaceholderNode, RollupNo use super::{AsLogicalPlan, LogicalExtensionCodec}; +/// Reconstruct cast target intent, using the legacy metadata and nullability +/// convention when `target_field` is absent. +fn cast_target_from_proto_parts( + data_type: DataType, + nullable: Option, + metadata: &std::collections::HashMap, + target_field: Option<&protobuf::Field>, +) -> std::result::Result { + if let Some(target_field) = target_field { + let field: Field = target_field.try_into()?; + if field.data_type() != &data_type { + return Err(proto_error(format!( + "Cast target field type {} does not match arrow_type {data_type}", + field.data_type() + ))); + } + return Ok(CastTarget::explicit(Arc::new(field))); + } + + if nullable.is_some_and(|nullable| !nullable) || !metadata.is_empty() { + Ok(CastTarget::explicit(Arc::new( + Field::new("", data_type, nullable.unwrap_or(true)) + .with_metadata(metadata.clone()), + ))) + } else { + Ok(CastTarget::type_only(data_type)) + } +} + /// Reconstruct a [`WriteOp`] from a [`protobuf::DmlNode`], reading the /// `merge_into` payload when the type tag is `MergeInto`. pub fn parse_write_op( @@ -426,11 +454,13 @@ pub fn parse_expr( "expr", codec, )?); - let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; - let field = data_type - .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); - Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field)))) + let field = cast_target_from_proto_parts( + cast.arrow_type.as_ref().required("arrow_type")?, + cast.nullable, + &cast.metadata, + cast.target_field.as_ref(), + )?; + Ok(Expr::Cast(Cast { expr, field })) } ExprType::TryCast(cast) => { let expr = Box::new(parse_required_expr( @@ -439,14 +469,13 @@ pub fn parse_expr( "expr", codec, )?); - let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; - let field = data_type - .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); - Ok(Expr::TryCast(TryCast::new_from_field( - expr, - Arc::new(field), - ))) + let field = cast_target_from_proto_parts( + cast.arrow_type.as_ref().required("arrow_type")?, + cast.nullable, + &cast.metadata, + cast.target_field.as_ref(), + )?; + Ok(Expr::TryCast(TryCast { expr, field })) } ExprType::Negative(negative) => Ok(Expr::Negative(Box::new( parse_required_expr(negative.expr.as_deref(), ctx, "expr", codec)?, diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 16c3468465541..13e21f5622c42 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -397,22 +397,40 @@ pub fn serialize_expr( } } Expr::Cast(Cast { expr, field }) => { + let target_field = field + .explicit_field() + .map(|field| field.as_ref().try_into()) + .transpose()?; let expr = Box::new(protobuf::CastNode { expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)), arrow_type: Some(field.data_type().try_into()?), - metadata: field.metadata().clone(), - nullable: Some(field.is_nullable()), + metadata: field.metadata().cloned().unwrap_or_default(), + nullable: Some( + field + .explicit_field() + .is_none_or(|field| field.is_nullable()), + ), + target_field, }); protobuf::LogicalExprNode { expr_type: Some(ExprType::Cast(expr)), } } Expr::TryCast(TryCast { expr, field }) => { + let target_field = field + .explicit_field() + .map(|field| field.as_ref().try_into()) + .transpose()?; let expr = Box::new(protobuf::TryCastNode { expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)), arrow_type: Some(field.data_type().try_into()?), - metadata: field.metadata().clone(), - nullable: Some(field.is_nullable()), + metadata: field.metadata().cloned().unwrap_or_default(), + nullable: Some( + field + .explicit_field() + .is_none_or(|field| field.is_nullable()), + ), + target_field, }); protobuf::LogicalExprNode { expr_type: Some(ExprType::TryCast(expr)), diff --git a/datafusion/proto/tests/cases/serialize.rs b/datafusion/proto/tests/cases/serialize.rs index 40548fb1cf335..0f7967e8cb8f7 100644 --- a/datafusion/proto/tests/cases/serialize.rs +++ b/datafusion/proto/tests/cases/serialize.rs @@ -24,8 +24,8 @@ use datafusion::execution::FunctionRegistry; use datafusion::prelude::SessionContext; use datafusion_common::ScalarValue; use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable, Placeholder}; +use datafusion_expr::{Cast, Expr, ScalarUDF, TryCast, Volatility}; use datafusion_expr::{ColumnarValue, HigherOrderUDF, col, create_udf, lambda, lit}; -use datafusion_expr::{Expr, ScalarUDF, Volatility}; use datafusion_functions::string; use datafusion_proto::bytes::{ Serializeable, logical_exprs_from_bytes_with_extension_codec, @@ -164,6 +164,73 @@ fn roundtrip_expr(expr: &Expr) -> Expr { Expr::from_bytes(&bytes).unwrap() } +#[test] +fn roundtrip_casts_preserve_type_only_and_explicit_targets() { + let explicit_field_without_metadata = Arc::new(Field::new("", DataType::Int64, true)); + let explicit_field_with_metadata = Arc::new( + Field::new("", DataType::Int64, true) + .with_metadata([("extension".to_string(), "custom".to_string())].into()), + ); + let input = Box::new(lit(1_i32)); + let cast_expressions = [ + Expr::Cast(Cast::new(input.clone(), DataType::Int64)), + Expr::Cast(Cast::new_from_field( + input.clone(), + Arc::clone(&explicit_field_without_metadata), + )), + Expr::Cast(Cast::new_from_field( + input.clone(), + Arc::clone(&explicit_field_with_metadata), + )), + Expr::TryCast(TryCast::new(input.clone(), DataType::Int64)), + Expr::TryCast(TryCast::new_from_field( + input.clone(), + explicit_field_without_metadata, + )), + Expr::TryCast(TryCast::new_from_field(input, explicit_field_with_metadata)), + ]; + + for expr in cast_expressions { + assert_eq!(expr, roundtrip_expr(&expr)); + } +} + +#[test] +fn decode_legacy_casts_without_target_field() { + let legacy_explicit_field = Arc::new( + Field::new("", DataType::Int64, false) + .with_metadata([("extension".to_string(), "custom".to_string())].into()), + ); + let input = Box::new(lit(1_i32)); + let cast_expressions = [ + Expr::Cast(Cast::new_from_field( + input.clone(), + Arc::clone(&legacy_explicit_field), + )), + Expr::TryCast(TryCast::new_from_field(input, legacy_explicit_field)), + ]; + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + + for expr in cast_expressions { + let mut proto = serialize_expr(&expr, &codec).unwrap(); + match proto.expr_type.as_mut().unwrap() { + datafusion_proto::protobuf::logical_expr_node::ExprType::Cast(cast) => { + cast.target_field = None; + } + datafusion_proto::protobuf::logical_expr_node::ExprType::TryCast( + try_cast, + ) => { + try_cast.target_field = None; + } + other => panic!("expected cast proto, got {other:?}"), + } + + let decoded = parse_expr(&proto, ctx.task_ctx().as_ref(), &codec).unwrap(); + assert_eq!(decoded, expr); + } +} + #[test] fn exact_roundtrip_linearized_binary_expr() { // (((A AND B) AND C) AND D) diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index dff18173ae32a..59c56e51dfadc 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -1217,11 +1217,7 @@ fn rewrite_expr_to_prunable( scalar_expr, schema, )?; - let left = Arc::new(phys_expr::CastExpr::new_with_target_field( - left, - Arc::clone(cast.target_field()), - None, - )); + let left = Arc::new(cast.with_new_expr(left)); // PruningPredicate does not support pruning on nested fields yet. // End-to-end nested-field pruning also requires Parquet statistics // extraction to agree with PruningPredicate on a stats representation @@ -1236,10 +1232,7 @@ fn rewrite_expr_to_prunable( scalar_expr, schema, )?; - let left = Arc::new(phys_expr::TryCastExpr::new( - left, - try_cast.cast_type().clone(), - )); + let left = Arc::new(try_cast.with_new_expr(left)); Ok((left, op, right)) } else if let Some(neg) = column_expr.downcast_ref::() { // `-col > lit()` --> `col < -lit()` diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index f1661028ed051..f2152bdae9723 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -378,14 +378,14 @@ impl SqlToRel<'_, S> { return not_impl_err!("CAST with format is not supported: {format}"); } - Ok(Expr::TryCast(TryCast::new_from_field( - Box::new(self.sql_expr_to_logical_expr( + Ok(Expr::TryCast(TryCast { + expr: Box::new(self.sql_expr_to_logical_expr( *expr, schema, planner_context, )?), - self.convert_data_type_to_field(&data_type)?, - ))) + field: self.convert_data_type_to_cast_target(&data_type)?, + })) } SQLExpr::TypedString(TypedString { @@ -397,10 +397,10 @@ impl SqlToRel<'_, S> { return plan_err!("Typed literal requires a string payload"); }; - Ok(Expr::Cast(Cast::new_from_field( - Box::new(lit(value)), - self.convert_data_type_to_field(&data_type)?, - ))) + Ok(Expr::Cast(Cast { + expr: Box::new(lit(value)), + field: self.convert_data_type_to_cast_target(&data_type)?, + })) } SQLExpr::IsNull(expr) => Ok(Expr::IsNull(Box::new( @@ -1124,12 +1124,12 @@ impl SqlToRel<'_, S> { return not_impl_err!("CAST with format is not supported: {format}"); } - let dt = self.convert_data_type_to_field(data_type)?; + let target = self.convert_data_type_to_cast_target(data_type)?; let expr = self.sql_expr_to_logical_expr(expr, schema, planner_context)?; // numeric constants are treated as seconds (rather as nanoseconds) // to align with postgres / duckdb semantics - let expr = match dt.data_type() { + let expr = match target.data_type() { DataType::Timestamp(TimeUnit::Nanosecond, tz) if expr.get_type(schema)? == DataType::Int64 => { @@ -1141,7 +1141,10 @@ impl SqlToRel<'_, S> { _ => expr, }; - Ok(Expr::Cast(Cast::new_from_field(Box::new(expr), dt))) + Ok(Expr::Cast(Cast { + expr: Box::new(expr), + field: target, + })) } /// Extracts the root expression and access chain from a compound expression. diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 1db11d66b7ec7..7f6f26e07b555 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -32,10 +32,10 @@ use datafusion_common::{ DFSchemaRef, Diagnostic, SchemaError, field_not_found, internal_err, plan_datafusion_err, }; -use datafusion_expr::Expr; use datafusion_expr::logical_plan::{LogicalPlan, LogicalPlanBuilder}; pub use datafusion_expr::planner::ContextProvider; use datafusion_expr::utils::find_column_exprs; +use datafusion_expr::{CastTarget, Expr}; use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo, TimezoneInfo}; use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption}; use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias}; @@ -668,14 +668,36 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { &self, sql_type: &SQLDataType, ) -> Result { - // First check if any of the registered type_planner can handle this type - if let Some(type_planner) = self.context_provider.get_type_planner() - && let Some(data_type) = type_planner.plan_type_field(sql_type)? - { - return Ok(data_type); + if let Some(field) = self.plan_custom_type_field(sql_type)? { + return Ok(field); } - // If no type_planner can handle this type, use the default conversion + self.convert_builtin_data_type_to_field(sql_type) + } + + pub(crate) fn convert_data_type_to_cast_target( + &self, + sql_type: &SQLDataType, + ) -> Result { + if let Some(field) = self.plan_custom_type_field(sql_type)? { + return Ok(CastTarget::explicit(field)); + } + + let field = self.convert_builtin_data_type_to_field(sql_type)?; + Ok(CastTarget::type_only(field.data_type().clone())) + } + + fn plan_custom_type_field(&self, sql_type: &SQLDataType) -> Result> { + match self.context_provider.get_type_planner() { + Some(type_planner) => type_planner.plan_type_field(sql_type), + None => Ok(None), + } + } + + fn convert_builtin_data_type_to_field( + &self, + sql_type: &SQLDataType, + ) -> Result { match sql_type { SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) => { // Arrays may be multi-dimensional. diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 898330018c708..50f74dd1a1fd2 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -262,7 +262,12 @@ impl Unparser<'_> { end_token: AttachedToken::empty(), }) } - Expr::Cast(Cast { expr, field }) => Ok(self.cast_to_sql(expr, field)?), + Expr::Cast(Cast { expr, field }) => { + let target = field.explicit_field().cloned().unwrap_or_else(|| { + field.data_type().clone().into_nullable_field_ref() + }); + Ok(self.cast_to_sql(expr, &target)?) + } Expr::Literal(value, _) => Ok(self.scalar_to_sql(value)?), Expr::Alias(Alias { expr, .. }) => self.expr_to_sql_inner(expr), Expr::WindowFunction(window_fun) => { @@ -556,10 +561,13 @@ impl Unparser<'_> { } Expr::TryCast(TryCast { expr, field }) => { let inner_expr = self.expr_to_sql_inner(expr)?; + let target = field.explicit_field().cloned().unwrap_or_else(|| { + field.data_type().clone().into_nullable_field_ref() + }); Ok(ast::Expr::Cast { kind: ast::CastKind::TryCast, expr: Box::new(inner_expr), - data_type: self.arrow_dtype_to_ast_dtype(field)?, + data_type: self.arrow_dtype_to_ast_dtype(&target)?, array: false, format: None, }) diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 00103bfd9f56a..c50e48c3d10e9 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -28,7 +28,7 @@ use arrow::datatypes::{TimeUnit::Nanosecond, *}; use common::MockContextProvider; use datafusion_common::{DFSchema, DataFusionError, Result, assert_contains}; use datafusion_expr::{ - ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs, + CastTarget, ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, Volatility, col, @@ -5343,6 +5343,43 @@ fn test_no_substring_registered_alt_syntax() { ); } +#[test] +fn test_builtin_sql_casts_use_type_only_targets() -> Result<()> { + let plan = logical_plan("SELECT CAST(1 AS SMALLINT)")?; + let LogicalPlan::Projection(projection) = plan else { + panic!("expected projection") + }; + let [Expr::Cast(cast)] = projection.expr.as_slice() else { + panic!("expected CAST expression") + }; + assert!(matches!(&cast.field, CastTarget::DataType(DataType::Int16))); + + let plan = logical_plan("SELECT TRY_CAST(1 AS SMALLINT)")?; + let LogicalPlan::Projection(projection) = plan else { + panic!("expected projection") + }; + let [Expr::TryCast(try_cast)] = projection.expr.as_slice() else { + panic!("expected TRY_CAST expression") + }; + assert!(matches!( + &try_cast.field, + CastTarget::DataType(DataType::Int16) + )); + + let plan = logical_plan("SELECT TIMESTAMP '2001-01-01 18:00:00'")?; + let LogicalPlan::Projection(projection) = plan else { + panic!("expected projection") + }; + let [Expr::Cast(cast)] = projection.expr.as_slice() else { + panic!("expected typed literal CAST expression") + }; + assert!(matches!( + &cast.field, + CastTarget::DataType(DataType::Timestamp(Nanosecond, None)) + )); + Ok(()) +} + #[test] fn test_custom_type_plan() -> Result<()> { let sql = "SELECT DATETIME '2001-01-01 18:00:00'"; @@ -5376,6 +5413,18 @@ fn test_custom_type_plan() -> Result<()> { } let plan = plan_sql(sql); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected projection") + }; + let Expr::Cast(cast) = &projection.expr[0] else { + panic!("expected CAST expression") + }; + assert!(matches!( + &cast.field, + CastTarget::Field(field) + if field.data_type() == &DataType::Timestamp(Nanosecond, None) + && field.metadata().is_empty() + )); assert_snapshot!( plan, @@ -5408,6 +5457,16 @@ fn test_custom_type_plan() -> Result<()> { ); let plan = plan_sql("SELECT UUID '00010203-0405-0607-0809-000102030506'"); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected projection") + }; + let Expr::Cast(cast) = &projection.expr[0] else { + panic!("expected CAST expression") + }; + assert!( + matches!(&cast.field, CastTarget::Field(field) if !field.metadata().is_empty()) + ); + assert_snapshot!( plan, @r#" @@ -5415,6 +5474,24 @@ fn test_custom_type_plan() -> Result<()> { EmptyRelation: rows=1 "# ); + + let plan = plan_sql("SELECT CAST(NULL AS ARRAY)"); + let LogicalPlan::Projection(projection) = &plan else { + panic!("expected projection") + }; + let Expr::Cast(cast) = &projection.expr[0] else { + panic!("expected CAST expression") + }; + let CastTarget::DataType(DataType::List(field)) = &cast.field else { + panic!("expected type-only list target") + }; + assert_eq!( + field + .metadata() + .get("ARROW:extension:name") + .map(String::as_str), + Some("arrow.uuid") + ); Ok(()) } diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs b/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs index 3dd62afe8f193..d914fb208112f 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/cast.rs @@ -16,7 +16,7 @@ // under the License. use crate::logical_plan::consumer::{ - SubstraitConsumer, field_from_substrait_type_without_names, + SubstraitConsumer, from_substrait_type_without_names, }; use datafusion::common::{DFSchema, substrait_err}; use datafusion::logical_expr::{Cast, Expr, TryCast}; @@ -38,13 +38,54 @@ pub async fn from_cast( ) .await?, ); - let field = field_from_substrait_type_without_names(consumer, output_type)?; + let data_type = from_substrait_type_without_names(consumer, output_type)?; if cast.failure_behavior() == ReturnNull { - Ok(Expr::TryCast(TryCast::new_from_field(input_expr, field))) + Ok(Expr::TryCast(TryCast::new(input_expr, data_type))) } else { - Ok(Expr::Cast(Cast::new_from_field(input_expr, field))) + Ok(Expr::Cast(Cast::new(input_expr, data_type))) } } None => substrait_err!("Cast expression without output type is not allowed"), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::logical_plan::consumer::utils::tests::test_consumer; + use datafusion::arrow::datatypes::DataType; + use datafusion::logical_expr::CastTarget; + use substrait::proto::expression::cast::FailureBehavior; + use substrait::proto::expression::{Literal, RexType, literal::LiteralType}; + use substrait::proto::r#type::{I64, Kind}; + use substrait::proto::{Expression, Type}; + + #[tokio::test] + async fn standard_casts_are_type_only() -> datafusion::common::Result<()> { + for failure_behavior in [FailureBehavior::ThrowException, ReturnNull] { + let cast = substrait_expression::Cast { + r#type: Some(Type { + kind: Some(Kind::I64(I64::default())), + }), + input: Some(Box::new(Expression { + rex_type: Some(RexType::Literal(Literal { + literal_type: Some(LiteralType::I32(1)), + ..Default::default() + })), + })), + failure_behavior: failure_behavior.into(), + }; + let consumer = test_consumer(); + let expr = from_cast(&consumer, &cast, &DFSchema::empty()).await?; + + let target = match expr { + Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { + field + } + expr => panic!("expected cast expression, got {expr}"), + }; + assert!(matches!(target, CastTarget::DataType(DataType::Int64))); + } + Ok(()) + } +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/cast.rs b/datafusion/substrait/src/logical_plan/producer/expr/cast.rs index 2a5a6fe5c3758..d79327be9c591 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/cast.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/cast.rs @@ -18,7 +18,7 @@ use crate::logical_plan::producer::{SubstraitProducer, to_substrait_type_from_field}; use crate::variation_const::DEFAULT_TYPE_VARIATION_REF; use datafusion::common::{DFSchemaRef, ScalarValue}; -use datafusion::logical_expr::{Cast, Expr, TryCast}; +use datafusion::logical_expr::{Cast, Expr, ExprSchemable, TryCast}; use substrait::proto::Expression; use substrait::proto::expression::cast::FailureBehavior; use substrait::proto::expression::literal::LiteralType; @@ -29,7 +29,8 @@ pub fn from_cast( cast: &Cast, schema: &DFSchemaRef, ) -> datafusion::common::Result { - let Cast { expr, field } = cast; + let Cast { expr, .. } = cast; + let (_, output_field) = Expr::Cast(cast.clone()).to_field(schema)?; // since substrait Null must be typed, so if we see a cast(null, dt), we make it a typed null if let Expr::Literal(lit, _) = expr.as_ref() { // only the untyped(a null scalar value) null literal need this special handling @@ -40,7 +41,8 @@ pub fn from_cast( nullable: true, type_variation_reference: DEFAULT_TYPE_VARIATION_REF, literal_type: Some(LiteralType::Null(to_substrait_type_from_field( - producer, field, + producer, + &output_field, )?)), }; return Ok(Expression { @@ -51,7 +53,7 @@ pub fn from_cast( Ok(Expression { rex_type: Some(RexType::Cast(Box::new( substrait::proto::expression::Cast { - r#type: Some(to_substrait_type_from_field(producer, field)?), + r#type: Some(to_substrait_type_from_field(producer, &output_field)?), input: Some(Box::new(producer.handle_expr(expr, schema)?)), failure_behavior: FailureBehavior::ThrowException.into(), }, @@ -64,11 +66,12 @@ pub fn from_try_cast( cast: &TryCast, schema: &DFSchemaRef, ) -> datafusion::common::Result { - let TryCast { expr, field } = cast; + let TryCast { expr, .. } = cast; + let (_, output_field) = Expr::TryCast(cast.clone()).to_field(schema)?; Ok(Expression { rex_type: Some(RexType::Cast(Box::new( substrait::proto::expression::Cast { - r#type: Some(to_substrait_type_from_field(producer, field)?), + r#type: Some(to_substrait_type_from_field(producer, &output_field)?), input: Some(Box::new(producer.handle_expr(expr, schema)?)), failure_behavior: FailureBehavior::ReturnNull.into(), }, @@ -85,7 +88,6 @@ mod tests { use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::common::DFSchema; use datafusion::execution::SessionStateBuilder; - use datafusion::logical_expr::ExprSchemable; use substrait::proto::expression_reference::ExprType; #[tokio::test]