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/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..0c61abde6661e 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -27,8 +27,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -161,8 +161,25 @@ impl ScalarUDFImpl for ArrowCastFunc { let [source_arg, type_arg] = take_function_args(self.name(), args)?; let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; - let new_expr = if source_type == target_type { - // the argument's data type is already the correct type + // `arrow_cast` names a `DataType` and nothing else: its declared return + // field (see `return_field_from_args`) never carries metadata. Dropping + // the cast is therefore only sound when the argument has nothing for the + // cast to drop - the cast target's metadata is authoritative, so a cast + // whose argument carries metadata is a metadata-changing cast even when + // the two data types are identical. + // + // Eliding it there would leave, for example, + // `arrow_cast(uuid_val, 'FixedSizeBinary(16)')` still labelled + // `ARROW:extension:name = arrow.uuid`, i.e. an extension value that + // escaped a cast back to its storage type. + // See https://github.com/apache/datafusion/issues/22079 + let is_noop = source_type == target_type && { + let (_, source_field) = source_arg.to_field(info.schema().as_ref())?; + source_field.metadata().is_empty() + }; + let new_expr = if is_noop { + // the argument's data type is already the correct type and it + // carries no metadata that the cast would strip source_arg } else { // Use an actual cast to get the correct type diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index cb3103d38c52a..cb7d1d9b02bcf 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -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, @@ -384,10 +381,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 +396,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 +1035,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![ diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..c624e3ffe5558 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -22,12 +22,13 @@ use std::sync::Arc; use crate::PhysicalExpr; use arrow::compute; use arrow::compute::CastOptions; -use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; use compute::can_cast_types; 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)] @@ -94,10 +95,16 @@ impl PhysicalExpr for TryCastExpr { } 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) + // `TryCastExpr` only knows a target `DataType`, so it stands in for a + // type-only cast target: the output keeps the source's name and + // nullability but not its metadata. Deriving that through the shared + // helper keeps it identical to what the logical `Expr::TryCast` reports. + let source_field = self.expr.return_field(input_schema)?; + Ok(cast_output_field( + &source_field, + &Field::new("", self.cast_type.clone(), true).into(), + false, + )) } fn children(&self) -> Vec<&Arc> { @@ -621,6 +628,30 @@ mod tests { .unwrap() } + #[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)]); 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/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..bd142cd311e9c 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -47,3 +47,56 @@ FROM ( 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); + +# 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 + +# `arrow_cast` names a storage type and nothing else, so casting an +# extension-typed value back to its own storage type is not a no-op: it strips +# the extension metadata. The same-type short circuit in `ArrowCastFunc::simplify` +# must not swallow it. +# See https://github.com/apache/datafusion/issues/22079 +query ?T +SELECT + arrow_cast(uuid_val, 'FixedSizeBinary(16)'), + arrow_metadata(arrow_cast(uuid_val, 'FixedSizeBinary(16)'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# ... and casting it to a different storage type, through `arrow_cast` rather +# than a SQL CAST, strips it too +query ?T +SELECT + arrow_cast(uuid_val, 'Binary'), + arrow_metadata(arrow_cast(uuid_val, 'Binary'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# An `arrow_cast` whose argument carries no metadata has nothing to strip and is +# still simplified away entirely +query TT +EXPLAIN SELECT arrow_cast(a, 'Int64') FROM (SELECT arrow_cast(1, 'Int64') AS a); +---- +logical_plan +01)Projection: Int64(1) AS arrow_cast(a,Utf8("Int64")) +02)--EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[1 as arrow_cast(a,Utf8("Int64"))] +02)--PlaceholderRowExec 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