diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 4e9839e2f7479..b53b4b5e01e76 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -22,14 +22,17 @@ use std::collections::HashSet; use std::fmt::Debug; use std::sync::Arc; -use crate::expr::{Alias, Sort, Unnest}; +use arrow::compute::can_cast_types; +use arrow::datatypes::Field; + +use crate::expr::{Alias, Cast, Sort, Unnest}; use crate::logical_plan::Projection; use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder}; use datafusion_common::TableReference; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{Column, DFSchema, Result}; +use datafusion_common::{Column, DFSchema, Result, plan_err}; mod guarantees; pub use guarantees::GuaranteeRewriter; @@ -255,11 +258,12 @@ 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(coerce_expr_to_field(*expr, dst_field, src_schema)? + .alias(name)) } #[expect(deprecated)] Expr::Wildcard { .. } => Ok(expr), @@ -270,9 +274,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(coerce_expr_to_field(expr, dst_field, src_schema)? + .alias(name)) } - _ => Ok(expr.cast_to(new_type, src_schema)?), + _ => coerce_expr_to_field(expr, dst_field, src_schema), } } } @@ -283,6 +288,36 @@ fn coerce_exprs_for_schema( .collect::>() } +/// Cast `expr` so that it produces `dst_field`, carrying that field's metadata on +/// the cast target. +/// +/// A cast's target metadata is authoritative, so coercing to a bare `DataType` +/// would produce an expression whose output field has no metadata, contradicting +/// the destination schema that the coercion is supposed to satisfy. +fn coerce_expr_to_field( + expr: Expr, + dst_field: &Field, + src_schema: &DFSchema, +) -> Result { + if dst_field.metadata().is_empty() { + return expr.cast_to(dst_field.data_type(), src_schema); + } + + let this_type = expr.get_type(src_schema)?; + if !can_cast_types(&this_type, dst_field.data_type()) { + return plan_err!( + "Cannot automatically convert {this_type} to {}", + dst_field.data_type() + ); + } + + let target = Arc::new( + Field::new("", dst_field.data_type().clone(), true) + .with_metadata(dst_field.metadata().clone()), + ); + Ok(Expr::Cast(Cast::new_from_field(Box::new(expr), target))) +} + /// 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 75b3a60af1465..5df4f29faba95 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -33,7 +33,6 @@ use crate::{ use arrow::compute::can_cast_types; use arrow::datatypes::FieldRef; use arrow::datatypes::{DataType, Field}; -use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::FieldExt; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, @@ -91,27 +90,14 @@ fn cast_output_field( target_field: &FieldRef, force_nullable: bool, ) -> Arc { - // Check if this is a "type-only" cast (target_field == DataType::X.into_nullable_field()) - let is_type_only = target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty(); - - let metadata = if is_type_only { - // Type-only cast: propagate source metadata, stripping extension type keys - let mut meta = source_field.metadata().clone(); - meta.remove(EXTENSION_TYPE_NAME_KEY); - meta.remove(EXTENSION_TYPE_METADATA_KEY); - meta - } else { - // Explicit target field: use target metadata exactly - target_field.metadata().clone() - }; - + // The cast target's metadata is authoritative: a cast produces the field its + // target describes. A type-only target describes a field with no metadata, so + // a plain `CAST(expr AS type)` produces no metadata. let mut f = source_field .as_ref() .clone() .with_data_type(target_field.data_type().clone()) - .with_metadata(metadata); + .with_metadata(target_field.metadata().clone()); if force_nullable { f = f.with_nullable(true); } @@ -1184,11 +1170,12 @@ mod tests { .with_data_type(DataType::Int32) .with_metadata(meta.clone()); - // col, alias, and cast should be metadata-preserving + // col and alias are metadata-preserving; a cast is not, because its + // target's metadata is authoritative and a type-only target carries none assert_eq!(meta, expr.metadata(&schema).unwrap()); assert_eq!(meta, expr.clone().alias("bar").metadata(&schema).unwrap()); assert_eq!( - meta, + FieldMetadata::from(HashMap::new()), expr.clone() .cast_to(&DataType::Int64, &schema) .unwrap() @@ -1515,10 +1502,9 @@ mod tests { .is_none(), "{cast_name}: Extension type name should be stripped when target has no extension metadata" ); - assert_eq!( - result_field.metadata().get("custom_key"), - Some(&"custom_value".to_string()), - "{cast_name}: Non-extension metadata should be preserved" + assert!( + result_field.metadata().get("custom_key").is_none(), + "{cast_name}: source metadata should not survive a type-only cast" ); if use_try_cast { assert!( diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 929a92abcca23..b85c8ae74af87 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -162,17 +162,14 @@ impl ScalarUDFImpl for ArrowCastFunc { let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; - // We can skip the cast only if: - // 1. The source and target types are the same - // 2. The source has no extension metadata that needs to be stripped + // We can skip the cast only if it would be a genuine no-op: the types + // already match and the source carries no metadata for the cast to drop. + // `arrow_cast`'s target is type-only, and a type-only target's (empty) + // metadata is authoritative, so any source metadata is metadata the cast + // removes. let new_expr = if source_type == target_type { - // Check if source has extension metadata let source_field = source_arg.to_field(info.schema())?; - let has_extension_metadata = source_field - .1 - .metadata() - .contains_key("ARROW:extension:name"); - if has_extension_metadata { + if !source_field.1.metadata().is_empty() { // Need to create a cast to strip extension metadata Expr::Cast(datafusion_expr::Cast { expr: Box::new(source_arg), diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index 565d396dad1b5..f7c599b5e4228 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -25,7 +25,6 @@ use crate::physical_expr::PhysicalExpr; use arrow::compute::{CastOptions, can_cast_types}; use arrow::datatypes::{DataType, DataType::*, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; -use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::nested_struct::{ @@ -232,19 +231,16 @@ impl CastExpr { )); } - // Type-only cast: pass through the source metadata and nullability, - // stripping extension type keys (the cast is to a plain storage type). + // Type-only cast: the target's metadata is authoritative and a type-only + // target carries none, so the output carries none. Nullability still + // follows the source, which a type-only target says nothing about. source_result.map(|source_field| { - let mut metadata = source_field.metadata().clone(); - metadata.remove(EXTENSION_TYPE_NAME_KEY); - metadata.remove(EXTENSION_TYPE_METADATA_KEY); - Arc::new( source_field .as_ref() .clone() .with_data_type(self.cast_type().clone()) - .with_metadata(metadata), + .with_metadata(self.target_field.metadata().clone()), ) }) } @@ -494,16 +490,13 @@ pub fn cast_with_target_field( && target_field.is_nullable() && target_field.metadata().is_empty(); - // For same-type casts, we can skip creating a CastExpr only if: - // 1. The target is type-only (no explicit metadata) - // 2. The source has no extension metadata that needs to be stripped - // Otherwise we need the CastExpr to strip extension metadata from the source. + // For same-type casts we can skip creating a CastExpr only if the cast would + // be a genuine no-op: the target is type-only, and the source carries no + // metadata for the cast to drop. Because the target's metadata is + // authoritative, a same-type cast is still meaningful when it clears metadata. if expr_type == *cast_type && is_type_only { let source_field = expr.return_field(input_schema)?; - let has_extension_metadata = source_field - .metadata() - .contains_key(EXTENSION_TYPE_NAME_KEY); - if !has_extension_metadata { + if source_field.metadata().is_empty() { return Ok(Arc::clone(&expr)); } } @@ -556,6 +549,7 @@ pub fn cast( #[cfg(test)] mod tests { use super::*; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use crate::expressions::column::col; @@ -1395,10 +1389,9 @@ mod tests { field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), "Type-only cast should strip extension type name from source" ); - assert_eq!( - field.metadata().get("custom_key"), - Some(&"custom_value".to_string()), - "Type-only cast should preserve non-extension metadata" + assert!( + field.metadata().get("custom_key").is_none(), + "Type-only cast should not carry source metadata" ); Ok(()) diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index c054026724fb1..0c80fd0fd12a2 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -25,7 +25,6 @@ use arrow::compute; use arrow::compute::CastOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; -use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use compute::can_cast_types; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; @@ -179,19 +178,16 @@ impl PhysicalExpr for TryCastExpr { )); } - // Pass-through metadata from source (stripping extension keys) + // A type-only target carries no metadata, and the target's metadata is + // authoritative, so the output carries none. source_result.map(|source_field| { - let mut metadata = source_field.metadata().clone(); - metadata.remove(EXTENSION_TYPE_NAME_KEY); - metadata.remove(EXTENSION_TYPE_METADATA_KEY); - Arc::new( source_field .as_ref() .clone() .with_data_type(self.cast_type().clone()) .with_nullable(true) // TRY_CAST is always nullable - .with_metadata(metadata), + .with_metadata(Default::default()), ) }) } @@ -310,16 +306,13 @@ pub fn try_cast_with_target_field( && target_field.is_nullable() && target_field.metadata().is_empty(); - // For same-type casts, we can skip creating a TryCastExpr only if: - // 1. The target is type-only (no explicit metadata) - // 2. The source has no extension metadata that needs to be stripped - // Otherwise we need the TryCastExpr to strip extension metadata from the source. + // For same-type casts we can skip creating a TryCastExpr only if the cast + // would be a genuine no-op: the target is type-only, and the source carries + // no metadata for the cast to drop. Because the target's metadata is + // authoritative, a same-type cast is still meaningful when it clears metadata. if expr_type == *cast_type && is_type_only { let source_field = expr.return_field(input_schema)?; - let has_extension_metadata = source_field - .metadata() - .contains_key(EXTENSION_TYPE_NAME_KEY); - if !has_extension_metadata { + if source_field.metadata().is_empty() { return Ok(Arc::clone(&expr)); } } @@ -355,6 +348,7 @@ mod tests { }, datatypes::*, }; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_physical_expr_common::physical_expr::fmt_sql; // runs an end-to-end test of physical type cast @@ -922,11 +916,9 @@ mod tests { field.metadata().get(EXTENSION_TYPE_METADATA_KEY).is_none(), "Type-only try_cast should strip extension type metadata" ); - // Non-extension metadata should pass through - assert_eq!( - field.metadata().get("custom_key"), - Some(&"custom_value".to_string()), - "Type-only try_cast should preserve non-extension metadata" + assert!( + field.metadata().get("custom_key").is_none(), + "Type-only try_cast should not carry source metadata" ); // Field name preserved, type changed, always nullable assert_eq!(field.name(), "a"); diff --git a/datafusion/sqllogictest/test_files/metadata.slt b/datafusion/sqllogictest/test_files/metadata.slt index 0fc74fa6cf602..4f169a15d3cd7 100644 --- a/datafusion/sqllogictest/test_files/metadata.slt +++ b/datafusion/sqllogictest/test_files/metadata.slt @@ -225,9 +225,9 @@ SELECT 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 query IT @@ -236,15 +236,15 @@ SELECT 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 @@ -380,9 +380,9 @@ SELECT 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 query IT @@ -391,9 +391,9 @@ SELECT 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) query IT @@ -402,15 +402,15 @@ SELECT 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