diff --git a/datafusion/expr-common/src/casts.rs b/datafusion/expr-common/src/casts.rs index 3518c02772672..292bcc285402f 100644 --- a/datafusion/expr-common/src/casts.rs +++ b/datafusion/expr-common/src/casts.rs @@ -23,10 +23,13 @@ use std::cmp::Ordering; +use std::sync::Arc; + use arrow::datatypes::{ - DataType, MAX_DECIMAL32_FOR_EACH_PRECISION, MAX_DECIMAL64_FOR_EACH_PRECISION, - MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL32_FOR_EACH_PRECISION, - MIN_DECIMAL64_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION, TimeUnit, + DataType, FieldRef, MAX_DECIMAL32_FOR_EACH_PRECISION, + MAX_DECIMAL64_FOR_EACH_PRECISION, MAX_DECIMAL128_FOR_EACH_PRECISION, + MIN_DECIMAL32_FOR_EACH_PRECISION, MIN_DECIMAL64_FOR_EACH_PRECISION, + MIN_DECIMAL128_FOR_EACH_PRECISION, TimeUnit, }; use arrow::temporal_conversions::{ MICROSECONDS, MILLISECONDS, MILLISECONDS_IN_DAY, NANOSECONDS, @@ -532,12 +535,158 @@ fn try_cast_binary( } } +/// Is `field` the synthesized, type-only target of a cast? +/// +/// `CAST(expr AS )` only names a [`DataType`], so the `FieldRef` that +/// carries it is synthesized with an empty name, no metadata, and the default +/// nullability. Such a field says nothing about the output beyond its data +/// type, and [`cast_output_field`] therefore takes the output's name and +/// nullability from the expression being cast instead. +/// +/// A target field that is *not* type-only came from somewhere that knows more +/// than the data type - a `TypePlanner` resolving an extension type, or a +/// schema the expression is being coerced to - and is authoritative. +pub fn is_type_only_cast_target(field: &FieldRef) -> bool { + field.name().is_empty() && field.is_nullable() && field.metadata().is_empty() +} + +/// Derive the output field of a cast from the field being cast and the cast's +/// target field. +/// +/// Both the logical (`Expr::Cast` / `Expr::TryCast`) and the physical +/// (`CastExpr` / `TryCastExpr`) representations of a cast pair a source field - +/// the field of the expression being cast - with a target field describing the +/// destination. This function is the single definition of how those two are +/// combined, so that the logical and physical layers cannot drift apart the way +/// they did in . +/// +/// The rule is: +/// +/// * the **data type** always comes from `target_field` +/// * the **metadata** always comes from `target_field`, *including* when it is +/// empty +/// * the **name** and **nullability** come from `target_field` when it carries +/// more than a data type (see [`is_type_only_cast_target`]), and from +/// `source_field` otherwise +/// +/// Metadata is never inherited from the source. A cast changes the storage type +/// of a value, while metadata such as `ARROW:extension:name` describes how to +/// interpret one particular storage type; carrying it across produces a field +/// claiming to be an extension type that it is not (see +/// ). A caller that wants the +/// source's metadata on the result has to say so, by putting it on +/// `target_field`. +/// +/// `force_nullable` is set for `TRY_CAST`, which yields `NULL` rather than an +/// error when the cast fails and so is always nullable. +pub fn cast_output_field( + source_field: &FieldRef, + target_field: &FieldRef, + force_nullable: bool, +) -> FieldRef { + let base = if is_type_only_cast_target(target_field) { + source_field + } else { + target_field + }; + let mut f = base + .as_ref() + .clone() + .with_data_type(target_field.data_type().clone()) + .with_metadata(target_field.metadata().clone()); + if force_nullable { + f = f.with_nullable(true); + } + Arc::new(f) +} + #[cfg(test)] mod tests { use super::*; + + fn field(name: &str, data_type: DataType, nullable: bool) -> FieldRef { + Arc::new(Field::new(name, data_type, nullable)) + } + + fn with_meta(f: FieldRef, key: &str, value: &str) -> FieldRef { + Arc::new( + f.as_ref() + .clone() + .with_metadata(HashMap::from([(key.to_string(), value.to_string())])), + ) + } + + #[test] + fn type_only_cast_target_is_recognised() { + assert!(is_type_only_cast_target(&field("", DataType::Int64, true))); + // a name, a non-default nullability, or metadata all make the target + // carry more than a data type + assert!(!is_type_only_cast_target(&field( + "uuid", + DataType::Int64, + true + ))); + assert!(!is_type_only_cast_target(&field( + "", + DataType::Int64, + false + ))); + assert!(!is_type_only_cast_target(&with_meta( + field("", DataType::Int64, true), + "k", + "v" + ))); + } + + #[test] + fn cast_output_field_does_not_inherit_source_metadata() { + let source = with_meta( + field("id", DataType::FixedSizeBinary(16), false), + "ARROW:extension:name", + "arrow.uuid", + ); + let target = field("", DataType::Binary, true); + + let out = cast_output_field(&source, &target, false); + + // the type-only target contributes the data type and (empty) metadata, + // the source contributes name and nullability + assert_eq!(out.data_type(), &DataType::Binary); + assert!(out.metadata().is_empty(), "{:?}", out.metadata()); + assert_eq!(out.name(), "id"); + assert!(!out.is_nullable()); + } + + #[test] + fn cast_output_field_takes_an_explicit_target_verbatim() { + let source = with_meta( + field("id", DataType::FixedSizeBinary(16), false), + "source_key", + "source_value", + ); + let target = with_meta( + field("uuid", DataType::FixedSizeBinary(16), true), + "ARROW:extension:name", + "arrow.uuid", + ); + + let out = cast_output_field(&source, &target, false); + + assert_eq!(out.as_ref(), target.as_ref()); + assert!(out.metadata().get("source_key").is_none()); + } + + #[test] + fn cast_output_field_force_nullable_is_for_try_cast() { + let source = field("id", DataType::Int32, false); + let target = field("", DataType::Int64, true); + + assert!(!cast_output_field(&source, &target, false).is_nullable()); + assert!(cast_output_field(&source, &target, true).is_nullable()); + } use arrow::compute::{CastOptions, cast_with_options}; use arrow::datatypes::{Field, Fields}; - use std::sync::Arc; + use std::collections::HashMap; #[derive(Debug, Clone)] enum ExpectedCast { diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 7a6ac3fc8b062..3087fa8000527 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -22,10 +22,12 @@ use std::collections::HashSet; use std::fmt::Debug; use std::sync::Arc; -use crate::expr::{Alias, Sort, Unnest}; +use crate::expr::{Alias, Cast, Sort, Unnest}; use crate::logical_plan::Projection; use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder}; +use arrow::datatypes::Field; + use datafusion_common::TableReference; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; @@ -255,11 +257,11 @@ fn coerce_exprs_for_schema( .into_iter() .enumerate() .map(|(idx, expr)| { - let new_type = dst_schema.field(idx).data_type(); - if new_type != &expr.get_type(src_schema)? { + let dst_field = dst_schema.field(idx); + if dst_field.data_type() != &expr.get_type(src_schema)? { match expr { Expr::Alias(Alias { expr, name, .. }) => { - Ok(expr.cast_to(new_type, src_schema)?.alias(name)) + Ok(cast_expr_to_field(*expr, dst_field, src_schema)?.alias(name)) } #[expect(deprecated)] Expr::Wildcard { .. } => Ok(expr), @@ -270,9 +272,10 @@ fn coerce_exprs_for_schema( // (see: https://github.com/apache/datafusion/issues/18818) Expr::Column(ref column) => { let name = column.name().to_owned(); - Ok(expr.cast_to(new_type, src_schema)?.alias(name)) + Ok(cast_expr_to_field(expr, dst_field, src_schema)? + .alias(name)) } - _ => Ok(expr.cast_to(new_type, src_schema)?), + _ => cast_expr_to_field(expr, dst_field, src_schema), } } } @@ -283,6 +286,44 @@ fn coerce_exprs_for_schema( .collect::>() } +/// Cast `expr` so that it matches `dst_field` - its data type *and* its metadata. +/// +/// [`ExprSchemable::cast_to`] only takes a `DataType`, so the cast it builds +/// carries a type-only target field. The cast target's metadata is +/// authoritative (see [`cast_output_field`]), so a type-only target would drop +/// whatever metadata `expr` had, leaving the coerced branch inconsistent with +/// the schema it was coerced *to* - which is precisely the metadata `dst_field` +/// advertises. Re-target the cast at `dst_field` instead. +/// +/// Only the destination's data type and metadata are taken; the name and +/// nullability stay those of `expr`, so that the cast target is still described +/// relative to the expression being cast. +/// +/// [`cast_output_field`]: datafusion_expr_common::casts::cast_output_field +fn cast_expr_to_field( + expr: Expr, + dst_field: &Field, + src_schema: &DFSchema, +) -> Result { + let (_, src_field) = expr.to_field(src_schema)?; + // Delegate to `cast_to` so that cast validity checks and the + // `ScalarSubquery` special case keep applying, then re-target the cast it + // produced at the destination field. + Ok(match expr.cast_to(dst_field.data_type(), src_schema)? { + Expr::Cast(Cast { expr, field }) => { + let target = Arc::new( + src_field + .as_ref() + .clone() + .with_data_type(field.data_type().clone()) + .with_metadata(dst_field.metadata().clone()), + ); + Expr::Cast(Cast::new_from_field(expr, target)) + } + other => other, + }) +} + /// Recursively un-alias an expressions #[inline] pub fn unalias(expr: Expr) -> Expr { diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 36b76f076d26a..c55604f15d339 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -38,6 +38,7 @@ use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, not_impl_err, plan_datafusion_err, plan_err, }; +use datafusion_expr_common::casts::cast_output_field; use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer; use datafusion_functions_window_common::field::WindowUDFFieldArgs; use std::sync::Arc; @@ -71,24 +72,6 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. -/// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. -fn cast_output_field( - source_field: &FieldRef, - target_type: &DataType, - force_nullable: bool, -) -> Arc { - let mut f = source_field - .as_ref() - .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); - if force_nullable { - f = f.with_nullable(true); - } - Arc::new(f) -} - fn scalar_arguments_for_fields( args: &[Expr], arg_fields: &[FieldRef], @@ -623,20 +606,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)), @@ -1158,11 +1137,13 @@ mod tests { .with_data_type(DataType::Int32) .with_metadata(meta.clone()); - // col, alias, and cast should be metadata-preserving + // col and alias should be metadata-preserving assert_eq!(meta, expr.metadata(&schema).unwrap()); assert_eq!(meta, expr.clone().alias("bar").metadata(&schema).unwrap()); + // a cast, on the other hand, is not: the cast target says what metadata + // the result carries, and a type-only cast target carries none assert_eq!( - meta, + FieldMetadata::default(), expr.clone() .cast_to(&DataType::Int64, &schema) .unwrap() @@ -1427,4 +1408,64 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_output_field_metadata() { + use crate::expr::{Cast, TryCast}; + + let source_meta = + HashMap::from([("source_key".to_string(), "source_value".to_string())]); + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_meta.clone())); + + // A target field carrying metadata is authoritative: the source metadata + // does not leak into the output. + let target_meta = + HashMap::from([("target_key".to_string(), "target_value".to_string())]); + let target = + Arc::new(Field::new("", DataType::Utf8, true).with_metadata(target_meta)); + + for expr in [ + Expr::Cast(Cast::new_from_field( + Box::new(col("foo")), + Arc::clone(&target), + )), + Expr::TryCast(TryCast::new_from_field( + Box::new(col("foo")), + Arc::clone(&target), + )), + ] { + let field = expr.to_field(&schema).unwrap().1; + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()) + ); + assert!(field.metadata().get("source_key").is_none()); + } + + // A target field with no metadata is authoritative too: the source's + // metadata describes `FixedSizeBinary(16)` and must not ride along onto + // the `Utf8` the cast produces. + // See https://github.com/apache/datafusion/issues/22079 + let bare = Arc::new(Field::new("", DataType::Utf8, true)); + for expr in [ + Expr::Cast(Cast::new_from_field( + Box::new(col("foo")), + Arc::clone(&bare), + )), + Expr::TryCast(TryCast::new_from_field( + Box::new(col("foo")), + Arc::clone(&bare), + )), + ] { + let field = expr.to_field(&schema).unwrap().1; + assert!( + field.metadata().is_empty(), + "expected no metadata, got {:?}", + field.metadata() + ); + } + assert!(!source_meta.is_empty()); + } } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index cb3103d38c52a..48d807cfa25d9 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use crate::physical_expr::PhysicalExpr; use arrow::compute::{CastOptions, can_cast_types}; -use arrow::datatypes::{DataType, DataType::*, FieldRef, Schema}; +use arrow::datatypes::{DataType, DataType::*, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; @@ -30,6 +30,7 @@ use datafusion_common::nested_struct::{ requires_nested_struct_cast, validate_data_type_compatibility, }; use datafusion_common::{Result, not_impl_err}; +use datafusion_expr_common::casts::{cast_output_field, is_type_only_cast_target}; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_expr_common::interval_arithmetic::Interval; use datafusion_expr_common::sort_properties::ExprProperties; @@ -150,19 +151,21 @@ impl CastExpr { &self.cast_options } + /// The field this cast produces. + /// + /// Delegates to the shared [`cast_output_field`], which is also what the + /// logical `Expr::Cast` uses to derive its output field, so the two layers + /// agree by construction. 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| { - Arc::new( - field - .as_ref() - .clone() - .with_data_type(self.cast_type().clone()), - ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + // An explicit target fully determines the output field, so do not resolve + // the child in that case: it may not be resolvable against this schema. + // `rewrite_file_row_index_expr` relies on this, wrapping a `Column` whose + // index is deliberately outside the schema the cast is asked about. + if !is_type_only_cast_target(&self.target_field) { + return Ok(Arc::clone(&self.target_field)); } + let source_field = self.expr.return_field(input_schema)?; + Ok(cast_output_field(&source_field, &self.target_field, false)) } /// Check if casting from the specified source type to the target type is a @@ -191,12 +194,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, @@ -313,12 +310,23 @@ impl PhysicalExpr for CastExpr { ) -> Result> { use datafusion_proto_models::protobuf; + // `arrow_type` stays populated for readers that predate + // `target_field`; `target_field` is only written when it carries more + // than the data type, so plans that do not use one encode byte for byte + // as they did before. + let target_field = if is_type_only_cast_target(&self.target_field) { + None + } else { + Some(self.target_field.as_ref().try_into()?) + }; + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( protobuf::PhysicalCastNode { expr: Some(Box::new(ctx.encode_child(self.expr())?)), arrow_type: Some(self.cast_type().try_into()?), + target_field, }, ))), })) @@ -358,7 +366,18 @@ impl CastExpr { internal_datafusion_err!("CastExpr is missing required field 'arrow_type'") })?; - Ok(Arc::new(CastExpr::new(expr, arrow_type.try_into()?, None))) + let target_field = match cast_expr.target_field.as_ref() { + Some(field) => Arc::new(Field::try_from(field)?), + // Encoded before `target_field` existed, or by a cast that only + // named a data type. + None => DataType::try_from(arrow_type)?.into_nullable_field_ref(), + }; + + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + target_field, + None, + ))) } } @@ -384,10 +403,13 @@ pub fn cast_with_options( /// preserving any explicit field semantics such as name, nullability, and /// metadata. /// -/// If the input expression already has the same data type, this helper still -/// preserves an explicit `target_field` by constructing a field-aware -/// [`CastExpr`]. Only the default synthesized field created by the legacy -/// type-only API is elided back to the original child expression. +/// The cast is elided only when it would be a genuine no-op, that is when the +/// field it produces (see [`cast_output_field`]) is already the field `expr` +/// has. Matching data types are not enough: a cast target's metadata is +/// authoritative, so `CAST(uuid_val AS FixedSizeBinary(16))` still *does* +/// something - it strips `ARROW:extension:name` - even though `uuid_val` is +/// already a `FixedSizeBinary(16)`. Eliding it there would leave the physical +/// plan reporting metadata that the logical plan has already dropped. pub fn cast_with_target_field( expr: Arc, input_schema: &Schema, @@ -396,8 +418,11 @@ 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)); + if expr_type == *cast_type { + let source_field = expr.return_field(input_schema)?; + if cast_output_field(&source_field, &target_field, false) == source_field { + return Ok(Arc::clone(&expr)); + } } let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { @@ -1032,6 +1057,77 @@ mod tests { Ok(()) } + #[test] + fn type_only_cast_does_not_inherit_source_metadata() -> Result<()> { + // The source's metadata describes the source's storage type; a cast + // produces a different one, so the metadata must not ride along. + // See https://github.com/apache/datafusion/issues/22079 + let metadata = HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(metadata.clone()), + ]); + + let expr = CastExpr::new(col("a", &schema)?, Binary, None); + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &Binary); + assert!(field.metadata().is_empty(), "{:?}", field.metadata()); + + // ... and the same holds for a cast that does not change the type: it + // is the only way to spell "drop this metadata". + let expr = CastExpr::new(col("a", &schema)?, FixedSizeBinary(16), None); + let field = expr.return_field(&schema)?; + assert!(field.metadata().is_empty(), "{:?}", field.metadata()); + + Ok(()) + } + + #[test] + fn same_type_cast_is_only_elided_when_it_is_a_no_op() -> Result<()> { + // Matching data types are not enough to drop a cast: the target's + // metadata is authoritative, so `CAST(uuid AS FixedSizeBinary(16))` + // still strips `ARROW:extension:name`. Eliding it would leave the + // physical plan reporting metadata the logical plan already dropped. + let metadata = HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), true).with_metadata(metadata), + Field::new("b", FixedSizeBinary(16), true), + ]); + + let kept = cast_with_target_field( + col("a", &schema)?, + &schema, + FixedSizeBinary(16).into_nullable_field_ref(), + None, + )?; + assert!( + kept.downcast_ref::().is_some(), + "a metadata-stripping cast must survive, got {kept}" + ); + assert!(kept.return_field(&schema)?.metadata().is_empty()); + + // `b` has nothing to strip, so the cast really is a no-op + let elided = cast_with_target_field( + col("b", &schema)?, + &schema, + FixedSizeBinary(16).into_nullable_field_ref(), + None, + )?; + assert!( + elided.downcast_ref::().is_none(), + "a no-op cast should be elided, got {elided}" + ); + + Ok(()) + } + #[test] fn struct_cast_validation_uses_nested_target_fields() -> Result<()> { let source_type = Struct(Fields::from(vec![ @@ -1258,6 +1354,7 @@ mod proto_tests { use datafusion_proto_models::protobuf::{ PhysicalCastNode, PhysicalExprNode, physical_expr_node, }; + use std::collections::HashMap; /// A `CastExpr` over an `Int32` column, casting to `Int64`. fn proto_cast_fixture() -> CastExpr { @@ -1277,11 +1374,83 @@ 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, + }, ))), } } + #[test] + fn target_field_survives_a_proto_round_trip() { + // `PhysicalCastNode` only carried a data type, so a cast to an + // extension type came back from serialization as a plain cast to the + // storage type - silently losing `ARROW:extension:name`. + let schema = Schema::new(vec![Field::new("a", Binary, true)]); + let target = Arc::new( + Field::new("uuid", FixedSizeBinary(16), true).with_metadata(HashMap::from([ + ("ARROW:extension:name".to_string(), "arrow.uuid".to_string()), + ])), + ); + let cast = CastExpr::new_with_target_field( + col("a", &schema).unwrap(), + Arc::clone(&target), + None, + ); + + let encoder = StubEncoder::ok(); + let node = cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .expect("CastExpr should encode to Some(node)"); + let encoded = match node.expr_type { + Some(physical_expr_node::ExprType::Cast(boxed)) => *boxed, + other => panic!("expected a CastExpr node, got {other:?}"), + }; + assert!( + encoded.target_field.is_some(), + "an explicit target field must be encoded" + ); + + let node = PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Cast(Box::new( + PhysicalCastNode { + expr: Some(Box::new(column_node("a"))), + ..encoded + }, + ))), + }; + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = CastExpr::try_from_proto(&node, &ctx).unwrap(); + let decoded = decoded + .downcast_ref::() + .expect("decoded expr should be a CastExpr"); + + assert_eq!(decoded.target_field(), &target); + } + + #[test] + fn a_type_only_target_field_is_not_encoded() { + let cast = proto_cast_fixture(); + let encoder = StubEncoder::ok(); + let node = cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + let encoded = match node.expr_type { + Some(physical_expr_node::ExprType::Cast(boxed)) => *boxed, + other => panic!("expected a CastExpr node, got {other:?}"), + }; + assert!( + encoded.target_field.is_none(), + "a type-only cast should encode exactly as it did before" + ); + } + #[test] fn try_to_proto_encodes_cast_expr() { 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..8bc17bd7eb363 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -65,3 +65,4 @@ pub use try_cast::{TryCastExpr, try_cast}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; +pub(crate) use try_cast::try_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..a57e811fefbf5 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -25,37 +25,60 @@ use arrow::compute::CastOptions; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow::record_batch::RecordBatch; use compute::can_cast_types; +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_common::casts::cast_output_field; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast #[derive(Debug, Eq)] pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + /// Field describing the desired output after casting + target_field: FieldRef, } // 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_field.eq(&other.target_field) } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.target_field.hash(state); } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. + /// + /// Synthesizes a type-only target field, which produces the long-standing + /// behaviour of `TRY_CAST(expr AS )`. Prefer + /// [`TryCastExpr::new_with_target_field`] when the destination is richer + /// than a `DataType` - for example an extension type resolved by a + /// `TypePlanner`, whose `ARROW:extension:name` lives on the field. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self::new_with_target_field(expr, cast_type.into_nullable_field_ref()) + } + + /// Create a new `TryCastExpr` with an explicit target `FieldRef`. + /// + /// The target field describes the destination of the cast: its data type + /// and its metadata are what the expression reports, per + /// [`cast_output_field`]. + /// + /// See [`TryCastExpr::new`] for the constructor that only accepts a + /// `DataType`. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { expr, target_field } } /// The expression to cast @@ -65,19 +88,24 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target_field.data_type() + } + + /// Field describing the output column after casting. + pub fn target_field(&self) -> &FieldRef { + &self.target_field } } 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 +118,14 @@ 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) + // Derived through the shared helper, which is also what the logical + // `Expr::TryCast` uses, so the two layers agree by construction. + let source_field = self.expr.return_field(input_schema)?; + Ok(cast_output_field(&source_field, &self.target_field, false)) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +136,16 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( + Ok(Arc::new(TryCastExpr::new_with_target_field( Arc::clone(&children[0]), - self.cast_type.clone(), + Arc::clone(&self.target_field), ))) } 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")] @@ -125,14 +153,25 @@ impl PhysicalExpr for TryCastExpr { &self, ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, ) -> Result> { + use datafusion_expr_common::casts::is_type_only_cast_target; use datafusion_proto_models::protobuf; + // `arrow_type` stays populated for readers that predate `target_field`; + // `target_field` is only written when it carries more than the data + // type, so plans that do not use one encode byte for byte as before. + let target_field = if is_type_only_cast_target(&self.target_field) { + None + } else { + Some(self.target_field.as_ref().try_into()?) + }; + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new( protobuf::PhysicalTryCastNode { expr: Some(Box::new(ctx.encode_child(&self.expr)?)), arrow_type: Some(self.cast_type().try_into()?), + target_field, }, ))), })) @@ -165,26 +204,62 @@ impl TryCastExpr { "TryCastExpr", "arrow_type", )?; - let cast_type: DataType = arrow_type.try_into()?; - Ok(Arc::new(TryCastExpr::new(expr, cast_type))) + let target_field = match try_cast.target_field.as_ref() { + Some(field) => Arc::new(arrow::datatypes::Field::try_from(field)?), + // Encoded before `target_field` existed, or by a cast that only + // named a data type. + None => DataType::try_from(arrow_type)?.into_nullable_field_ref(), + }; + + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + target_field, + ))) } } /// Return a PhysicalExpression representing `expr` casted to /// `cast_type`, if any casting is needed. /// -/// Note that such casts may lose type information +/// Note that such casts may lose type information: the crate-internal +/// `try_cast_with_target_field` is used when the destination is described by a +/// `Field` rather than a bare `DataType`. pub fn try_cast( expr: Arc, input_schema: &Schema, cast_type: DataType, +) -> Result> { + try_cast_with_target_field(expr, input_schema, cast_type.into_nullable_field_ref()) +} + +/// Return a PhysicalExpression representing `expr` `TRY_CAST`ed to +/// `target_field`, if any casting is needed. +/// +/// Mirrors [`cast_with_target_field`]: the cast is elided only when it would be +/// a genuine no-op, that is when the field it produces is already the field +/// `expr` has. Matching data types alone are not enough, because the target's +/// metadata is authoritative and a same-type cast is how metadata is dropped. +/// +/// [`cast_with_target_field`]: crate::expressions::cast_with_target_field +pub(crate) fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: FieldRef, ) -> Result> { let expr_type = expr.data_type(input_schema)?; - if expr_type == cast_type { - Ok(Arc::clone(&expr)) - } else if can_cast_types(&expr_type, &cast_type) { - Ok(Arc::new(TryCastExpr::new(expr, cast_type))) + let cast_type = target_field.data_type(); + if expr_type == *cast_type { + let source_field = expr.return_field(input_schema)?; + if cast_output_field(&source_field, &target_field, false) == source_field { + return Ok(Arc::clone(&expr)); + } + } + 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}") } @@ -621,6 +696,84 @@ mod tests { .unwrap() } + #[test] + fn try_cast_with_target_field_carries_target_metadata() -> Result<()> { + // A `TRY_CAST` to an extension type must report the target's metadata, + // just like `CAST`. Before `TryCastExpr` had a target field there was + // nowhere to put it and the planner rejected the query outright. + let schema = Schema::new(vec![Field::new("a", DataType::Binary, true)]); + let target = Arc::new( + Field::new("uuid", DataType::FixedSizeBinary(16), true).with_metadata( + std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]), + ), + ); + + let expr = + try_cast_with_target_field(col("a", &schema)?, &schema, Arc::clone(&target))?; + + assert_eq!(expr.return_field(&schema)?, target); + assert_eq!(expr.data_type(&schema)?, DataType::FixedSizeBinary(16)); + Ok(()) + } + + #[test] + fn same_type_try_cast_is_only_elided_when_it_is_a_no_op() -> Result<()> { + let metadata = std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), true) + .with_metadata(metadata.clone()), + Field::new("b", DataType::FixedSizeBinary(16), true), + ]); + + // stripping `a`'s metadata is real work, so the cast has to stay + let kept = try_cast(col("a", &schema)?, &schema, DataType::FixedSizeBinary(16))?; + assert!( + kept.downcast_ref::().is_some(), + "a metadata-stripping TRY_CAST must survive, got {kept}" + ); + assert!(kept.return_field(&schema)?.metadata().is_empty()); + + // `b` has nothing to strip + let elided = + try_cast(col("b", &schema)?, &schema, DataType::FixedSizeBinary(16))?; + assert!( + elided.downcast_ref::().is_none(), + "a no-op TRY_CAST should be elided, got {elided}" + ); + + Ok(()) + } + + #[test] + fn try_cast_does_not_inherit_source_metadata() -> Result<()> { + // `TRY_CAST` follows the same rule as `CAST`: the source's metadata + // describes the source's storage type and is not carried across. + // See https://github.com/apache/datafusion/issues/22079 + let metadata = std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), false) + .with_metadata(metadata.clone()), + ]); + + let expr = try_cast(col("a", &schema)?, &schema, DataType::Binary)?; + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Binary); + assert!(field.metadata().is_empty(), "{:?}", field.metadata()); + + Ok(()) + } + #[test] fn test_fmt_sql() -> Result<()> { let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); @@ -676,7 +829,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, + }, ))), } } @@ -733,6 +890,75 @@ mod proto_tests { assert!(try_cast.expr().downcast_ref::().is_some()); } + #[test] + fn target_field_survives_a_proto_round_trip() { + let schema = Schema::new(vec![Field::new("a", DataType::Binary, true)]); + let target = Arc::new( + Field::new("uuid", DataType::FixedSizeBinary(16), true).with_metadata( + std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]), + ), + ); + let try_cast = TryCastExpr::new_with_target_field( + col("a", &schema).unwrap(), + Arc::clone(&target), + ); + + let encoder = StubEncoder::ok(); + let node = try_cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .expect("TryCastExpr should encode to Some(node)"); + + // rebuild the node with a decodable child, then decode it + let encoded = match node.expr_type { + Some(physical_expr_node::ExprType::TryCast(boxed)) => *boxed, + other => panic!("expected a TryCastExpr node, got {other:?}"), + }; + assert!( + encoded.target_field.is_some(), + "an explicit target field must be encoded" + ); + let node = PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::TryCast(Box::new( + PhysicalTryCastNode { + expr: Some(Box::new(column_node("a"))), + ..encoded + }, + ))), + }; + + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = TryCastExpr::try_from_proto(&node, &ctx).unwrap(); + let decoded = decoded + .downcast_ref::() + .expect("decoded expr should be a TryCastExpr"); + + assert_eq!(decoded.target_field(), &target); + } + + #[test] + fn a_type_only_target_field_is_not_encoded() { + let try_cast = try_cast_fixture(); + let encoder = StubEncoder::ok(); + let node = try_cast + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + let encoded = match node.expr_type { + Some(physical_expr_node::ExprType::TryCast(boxed)) => *boxed, + other => panic!("expected a TryCastExpr node, got {other:?}"), + }; + assert!( + encoded.target_field.is_none(), + "a type-only cast should encode exactly as it did before" + ); + } + #[test] fn try_from_proto_rejects_non_try_cast_node() { let node = column_node("a"); diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 9a0bdc33da8e9..0874340ff300e 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -27,7 +27,7 @@ use crate::{ use arrow::datatypes::Schema; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::FieldExt; -use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata}; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DFSchema, Result, ScalarValue, TableReference, ToDFSchema, exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, @@ -391,19 +391,7 @@ pub fn create_physical_expr( None, ), Expr::TryCast(TryCast { expr, field }) => { - if !field.metadata().is_empty() { - let (_, src_field) = expr.to_field(input_dfschema)?; - return plan_err!( - "TryCast from {} to {} is not supported", - format_type_and_metadata( - src_field.data_type(), - Some(src_field.metadata()), - ), - format_type_and_metadata(field.data_type(), Some(field.metadata())) - ); - } - - expressions::try_cast( + expressions::try_cast_with_target_field( create_physical_expr( expr, input_dfschema, @@ -411,7 +399,7 @@ pub fn create_physical_expr( planning_ctx, )?, input_schema, - field.data_type().clone(), + Arc::clone(field), ) } Expr::Not(expr) => expressions::not(create_physical_expr( diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2a24eb60e6fbc..4b7236a9ddcc5 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -270,6 +270,22 @@ impl ProjectionExec { )) } + /// Returns whether this projection's output metadata differs from the + /// metadata derived from its expressions and input schema. + fn overrides_metadata(&self) -> Result { + let derived_schema = self + .projector + .projection() + .project_schema(self.input.schema().as_ref())?; + let output_schema = self.schema(); + Ok(derived_schema.metadata() != output_schema.metadata() + || derived_schema + .fields() + .iter() + .zip(output_schema.fields()) + .any(|(derived, output)| derived.metadata() != output.metadata())) + } + /// Collect reverse alias mapping from projection expressions. /// The result hash map is a map from aliased Column in parent to original expr. fn collect_reverse_alias( @@ -1014,6 +1030,10 @@ pub fn remove_unnecessary_projections( plan: Arc, ) -> Result>> { let maybe_modified = if let Some(projection) = plan.downcast_ref::() { + // Removing a projection with observable metadata can change query results. + if projection.overrides_metadata()? { + return Ok(Transformed::no(plan)); + } // If the projection does not cause any change on the input, we can // safely remove it: if is_projection_removable(projection) { @@ -1031,6 +1051,7 @@ pub fn remove_unnecessary_projections( /// Compare the inputs and outputs of the projection. All expressions must be /// columns without alias, and projection does not change the order of fields. +/// The input and output schemas must also match exactly to preserve metadata. /// For example, if the input schema is `a, b`, `SELECT a, b` is removable, /// but `SELECT b, a` and `SELECT a+1, b` and `SELECT a AS c, b` are not. fn is_projection_removable(projection: &ProjectionExec) -> bool { @@ -1041,6 +1062,7 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool { }; col.name() == proj_expr.alias && col.index() == idx }) && exprs.len() == projection.input().schema().fields().len() + && projection.schema() == projection.input().schema() } /// Given the expression set of a projection, checks if the projection causes @@ -1074,13 +1096,17 @@ pub fn new_projections_for_columns( } /// Creates a new [`ProjectionExec`] instance with the given child plan and -/// projected expressions. +/// projected expressions, preserving the original output metadata. pub fn make_with_child( projection: &ProjectionExec, child: &Arc, ) -> Result> { - ProjectionExec::try_new(projection.expr().to_vec(), Arc::clone(child)) - .map(|e| Arc::new(e) as _) + ProjectionExec::try_new_with_schema_metadata( + projection.expr().to_vec(), + Arc::clone(child), + projection.schema().as_ref(), + ) + .map(|e| Arc::new(e) as _) } /// Returns `true` if all the expressions in the argument are `Column`s. @@ -1331,12 +1357,20 @@ pub fn update_join_filter( fn try_collapse_projection_chain( outer: &ProjectionExec, ) -> Result>> { + if outer.overrides_metadata()? { + return Ok(None); + } + let mut current_exprs: Vec = outer.expr().to_vec(); let mut current_input: Arc = Arc::clone(outer.input()); let mut column_ref_map: HashMap = HashMap::new(); let mut collapsed_any = false; 'outer: while let Some(inner_proj) = current_input.downcast_ref::() { + if inner_proj.overrides_metadata()? { + break; + } + // Collect the column references usage in the outer projection. column_ref_map.clear(); for proj_expr in ¤t_exprs { @@ -1386,8 +1420,13 @@ fn try_collapse_projection_chain( } // To unify 3 or more sequential projections: + // Preserve the outer projection's output metadata. let unified: Arc = - Arc::new(ProjectionExec::try_new(current_exprs, current_input)?); + Arc::new(ProjectionExec::try_new_with_schema_metadata( + current_exprs, + current_input, + outer.schema().as_ref(), + )?); remove_unnecessary_projections(unified).data().map(Some) } @@ -1517,11 +1556,14 @@ mod tests { use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::StringArray; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; use datafusion_common::stats::{ColumnStatistics, Precision, Statistics}; - use datafusion_expr::Operator; + use datafusion_expr::{Operator, ScalarUDF}; + use datafusion_functions::core::arrow_metadata::ArrowMetadataFunc; + use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::expressions::{ BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; @@ -1566,6 +1608,91 @@ mod tests { Ok(()) } + fn identity_projection_with_metadata( + input: Arc, + ) -> Result> { + let metadata_schema = + Schema::new_with_metadata( + vec![Field::new("i", DataType::Int32, true).with_metadata( + HashMap::from([("event_field".to_string(), "true".to_string())]), + )], + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + ); + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }], + input, + &metadata_schema, + )?)) + } + + #[test] + fn test_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata(test::scan_partitioned(1))?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_make_with_child_preserves_output_metadata() -> Result<()> { + let projection = identity_projection_with_metadata(test::scan_partitioned(1))?; + let projection = projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + + let rebuilt = make_with_child(projection, &test::scan_partitioned(1))?; + + assert_eq!(rebuilt.schema(), projection.schema()); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_parent_blocks_projection_collapse() -> Result<()> { + let inner = identity_projection_with_metadata(test::scan_partitioned(1))?; + let arrow_metadata = ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + ); + let outer: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(arrow_metadata), + alias: "metadata".to_string(), + }], + inner, + )?); + + let outer_projection = outer + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + assert!(try_collapse_projection_chain(outer_projection)?.is_none()); + + let optimized = remove_unnecessary_projections(outer)?.data; + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("metadata expression should return Utf8"); + assert_eq!(values.value(0), "true"); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new( diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index e685b947bbaaf..fb3ab2267ddfd 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1187,11 +1187,19 @@ message PhysicalCaseNode { message PhysicalTryCastNode { PhysicalExprNode expr = 1; datafusion_common.ArrowType arrow_type = 2; + // The full cast target. A cast target may describe more than a data type - + // an extension type carries `ARROW:extension:name` in its field metadata - + // and that metadata is what the cast's output reports. Optional so that + // plans encoded before this field existed still decode, falling back to + // `arrow_type`. + optional datafusion_common.Field target_field = 3; } message PhysicalCastNode { PhysicalExprNode expr = 1; datafusion_common.ArrowType arrow_type = 2; + // See PhysicalTryCastNode.target_field. + optional 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 4fcca074c8199..e28f4965edf31 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -18088,6 +18088,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)?; @@ -18095,6 +18098,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() } } @@ -18108,12 +18114,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 @@ -18137,6 +18146,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)), } } @@ -18158,6 +18168,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 => { @@ -18172,11 +18183,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__, }) } } @@ -22052,6 +22070,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)?; @@ -22059,6 +22080,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() } } @@ -22072,12 +22096,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 @@ -22101,6 +22128,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)), } } @@ -22122,6 +22150,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 => { @@ -22136,11 +22165,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__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index d14195ecc54a0..0918aae99d970 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1846,6 +1846,13 @@ pub struct PhysicalTryCastNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, tag = "2")] pub arrow_type: ::core::option::Option, + /// The full cast target. A cast target may describe more than a data type - + /// an extension type carries `ARROW:extension:name` in its field metadata - + /// and that metadata is what the cast's output reports. Optional so that + /// plans encoded before this field existed still decode, falling back to + /// `arrow_type`. + #[prost(message, optional, tag = "3")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalCastNode { @@ -1853,6 +1860,9 @@ pub struct PhysicalCastNode { pub expr: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, tag = "2")] pub arrow_type: ::core::option::Option, + /// See PhysicalTryCastNode.target_field. + #[prost(message, optional, tag = "3")] + pub target_field: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalNegativeNode { diff --git a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..ccecd67b9150a 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -45,5 +45,59 @@ FROM ( ---- 00010203040506070809000102030506 arrow.uuid -statement error DataFusion error: Optimizer rule 'simplify_expressions' failed[\s\S]*TryCast from FixedSizeBinary\(16\) to FixedSizeBinary\(16\)<\{"ARROW:extension:name": "arrow\.uuid"\}> is not supported -SELECT TRY_CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID); +# TRY_CAST to an extension type preserves the target's extension metadata, the +# same as CAST. This used to be a planning error, because the physical +# `TryCastExpr` had nowhere to put the target field. +query ?T +SELECT + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + arrow_metadata( + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + 'ARROW:extension:name' + ); +---- +00010203040506070809000102030506 arrow.uuid + +# ... including when the value comes from a column rather than a literal +query ?T +SELECT + TRY_CAST(raw AS UUID), + arrow_metadata(TRY_CAST(raw AS UUID), 'ARROW:extension:name') +FROM ( + VALUES ( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + ) +) AS uuids(raw); +---- +00010203040506070809000102030506 arrow.uuid + +# A TRY_CAST that only names a data type still drops the source's metadata +query ?T +SELECT + TRY_CAST(uuid_val AS BYTEA), + arrow_metadata(TRY_CAST(uuid_val AS BYTEA), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# Casting an extension-typed value to a different storage type drops the +# extension metadata: `arrow.uuid` describes a `FixedSizeBinary(16)`, and a +# `Binary` column that claims to be one is simply wrong. +# See https://github.com/apache/datafusion/issues/22079 +query ?T +SELECT + CAST(uuid_val AS BYTEA), + arrow_metadata(CAST(uuid_val AS BYTEA), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL diff --git a/datafusion/sqllogictest/test_files/metadata.slt b/datafusion/sqllogictest/test_files/metadata.slt index 0fc74fa6cf602..41c6f8166e7bf 100644 --- a/datafusion/sqllogictest/test_files/metadata.slt +++ b/datafusion/sqllogictest/test_files/metadata.slt @@ -167,6 +167,24 @@ no_baz no_foo NULL +# Regression test: a UNION branch that has to be coerced to the union's output +# type keeps the metadata that output schema advertises. The coercion inserts a +# CAST, and a cast target's metadata is authoritative, so the cast has to be +# built against the destination *field* rather than just its data type - +# otherwise the coerced branch drops metadata the Union schema still claims and +# the physical plan no longer matches the logical one. +query TT +select name, arrow_metadata(name, 'metadata_key') from ( + SELECT nonnull_name as name FROM "table_with_metadata" + UNION ALL + SELECT NULL::string as name +) order by name; +---- +no_bar the nonnull_name field +no_baz the nonnull_name field +no_foo the nonnull_name field +NULL the nonnull_name field + # Regression test: missing schema metadata from union when schema with metadata isn't the first one # and also ensure it works fine with multiple unions query T @@ -218,33 +236,35 @@ FROM table_with_metadata; 2020-09-08 2020-09-08 -# Regression test: CAST should preserve source field metadata +# A CAST does not inherit the source field's metadata: the metadata describes +# one particular storage type, and the cast produces a different one. +# See https://github.com/apache/datafusion/issues/22079 query DT SELECT CAST(ts AS DATE) as casted, arrow_metadata(CAST(ts AS DATE), 'metadata_key') FROM table_with_metadata; ---- -2020-09-08 ts non-nullable field -2020-09-08 ts non-nullable field -2020-09-08 ts non-nullable field +2020-09-08 NULL +2020-09-08 NULL +2020-09-08 NULL -# Regression test: CAST preserves metadata on integer column +# ... and that holds for a widening cast between integer types too query IT SELECT CAST(id AS BIGINT) as casted, arrow_metadata(CAST(id AS BIGINT), 'metadata_key') FROM table_with_metadata; ---- -1 the id field -NULL the id field -3 the id field +1 NULL +NULL NULL +3 NULL # Regression test: CAST with single-argument arrow_metadata (returns full map) query ? select arrow_metadata(CAST(id AS BIGINT)) from table_with_metadata limit 1; ---- -{metadata_key: the id field} +{} # Regression test: distinct with cast query D @@ -373,44 +393,45 @@ select arrow_metadata(id) from table_with_metadata limit 1; ---- {metadata_key: the id field} -# Regression test: TRY_CAST should preserve source field metadata +# TRY_CAST follows the same rule as CAST: the source's metadata is not +# inherited by the result. query DT SELECT TRY_CAST(ts AS DATE) as try_casted, arrow_metadata(TRY_CAST(ts AS DATE), 'metadata_key') FROM table_with_metadata; ---- -2020-09-08 ts non-nullable field -2020-09-08 ts non-nullable field -2020-09-08 ts non-nullable field +2020-09-08 NULL +2020-09-08 NULL +2020-09-08 NULL -# Regression test: TRY_CAST preserves metadata on integer column +# ... including on an integer column query IT SELECT TRY_CAST(id AS BIGINT) as try_casted, arrow_metadata(TRY_CAST(id AS BIGINT), 'metadata_key') FROM table_with_metadata; ---- -1 the id field -NULL the id field -3 the id field +1 NULL +NULL NULL +3 NULL -# Regression test: TRY_CAST preserves metadata even when cast fails (returns NULL) +# ... and when the cast fails and the value itself is NULL query IT SELECT TRY_CAST(name AS INT) as try_casted, arrow_metadata(TRY_CAST(name AS INT), 'metadata_key') FROM table_with_metadata; ---- -NULL the name field -NULL the name field -NULL the name field +NULL NULL +NULL NULL +NULL NULL # Regression test: TRY_CAST with single-argument arrow_metadata (returns full map) query ? select arrow_metadata(TRY_CAST(id AS BIGINT)) from table_with_metadata limit 1; ---- -{metadata_key: the id field} +{} # with_metadata: attach a single key and read it back query T