From a75e52fa3eced87519039a064b761548f8a06a53 Mon Sep 17 00:00:00 2001 From: Dewey Dunnington Date: Tue, 1 Sep 2026 19:31:44 +0000 Subject: [PATCH] Align metadata propagation through Physical and Logical casts (#23169) - Closes #22079 - Closes #24724 The logical `Expr::Cast` and `Expr::TryCast` have a `FieldRef` target that was added in https://github.com/apache/datafusion/pull/18136 so that logical casts can express a cast to an extension type. In combination with a SQL type planner ( https://github.com/apache/datafusion/pull/20676 ) and an optimizer rule, this enabled casting to/from extension types with custom semantics to actually occur. The ability to do this was reverted by https://github.com/apache/datafusion/pull/20836 (which removed the original test) and I am not sure that ability ever made it into a release. When investigating this issue, it became clear the logical and physical cast behaviour had diverged with respect to the target field. This PR strips specific metadata keys (extension name and extension metadata) when propagating metadata from the source of a cast to the target (because doing so may result in an invalid destination field that consumers could reject), and propagates all metadata from the (logical) cast target field (e.g., so that a cast to an extension type represented by the cast target field will have a `to_field()` that communicates the extension type). For the physical cast, this behaviour is replicated exactly (I hope). Note that actually casting to an extension type can be implemented with an optimizer rule, planner, or by the mechanism I have in the works in https://github.com/apache/datafusion/pull/21071 . Yes It was in practice not common to create a `Expr::Cast` with field metadata internally and thus I don't think users will see metadata changes from the inclusion of metadata from the target field. I would be surprised if stripping the extension name/metadata from the source was disruptive (it was more likely to have caused errors). Superceeds an earlier but similar attempt ( https://github.com/apache/datafusion/pull/22162 ). --------- Co-authored-by: Andrew Lamb Co-authored-by: Tim Saucer (cherry picked from commit 124291e113b8a0991a3abad585703c4a55b948cf) --- Cargo.lock | 1 + .../custom_data_source/custom_file_casts.rs | 5 +- datafusion/expr/src/expr_schema.rs | 167 +++++++- datafusion/functions/src/core/arrow_cast.rs | 26 +- .../functions/src/core/arrow_try_cast.rs | 24 +- .../physical-expr-adapter/src/rewrite.rs | 7 +- datafusion/physical-expr/Cargo.toml | 1 + .../physical-expr/src/expressions/cast.rs | 391 ++++++++++++++---- .../physical-expr/src/expressions/mod.rs | 2 +- .../physical-expr/src/expressions/try_cast.rs | 354 +++++++++++++++- datafusion/physical-expr/src/planner.rs | 161 ++++++-- .../proto/src/logical_plan/from_proto.rs | 6 +- .../tests/cases/roundtrip_logical_plan.rs | 26 ++ .../cast_extension_type_metadata.slt | 54 ++- 14 files changed, 1070 insertions(+), 155 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c9f916386df71..2483c2f325843 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2406,6 +2406,7 @@ name = "datafusion-physical-expr" version = "55.0.0" dependencies = [ "arrow", + "arrow-schema", "criterion", "datafusion-common", "datafusion-expr", diff --git a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs index 71addc6d1bcb0..202c0a71257e9 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs @@ -188,10 +188,11 @@ impl PhysicalExprAdapter for CustomCastsPhysicalExprAdapter { if let Some(cast) = expr.downcast_ref::() { let input_data_type = cast.expr().data_type(&self.physical_file_schema)?; - let output_data_type = cast.target_field().data_type(); + let output_field = cast.target_field(); if !cast.is_bigger_cast(&input_data_type) { return not_impl_err!( - "Unsupported CAST from {input_data_type} to {output_data_type}" + "Unsupported CAST from {input_data_type} to {}", + output_field.data_type() ); } } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 8927fcf4d0bbe..cdd66f99ad3b7 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -31,6 +31,7 @@ use crate::{LogicalPlan, Projection, Subquery, WindowFunctionDefinition, utils}; 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, @@ -69,18 +70,46 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. +/// Derives the output field for a cast expression from the source and target fields. +/// +/// Metadata handling: +/// - Type-only casts (i.e., target_field == DataType::SomeDataType.into_nullable_field()) +/// propagate non extension-type metadata from the source. This is for backward compatibility +/// (casts have propagated source metadata for many if not all previous versions), recognizing +/// that the return type of `::` should have the +/// return type of `` (e.g., casting arrow.json to utf8). +/// - All other casts preserve target metadata exactly. This ensures in particular that output +/// metadata when casting to an extension type contains the extension information in the +/// output field. Callers that wish to have some mix of source and target metadata can use +/// Alias or construct an output field themselves (whose metadata will be used directly). +/// /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + 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() + }; + let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target_field.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -462,7 +491,8 @@ impl ExprSchemable for Expr { /// - **Aliases**: Merge underlying expr metadata with alias-specific metadata, preferring the alias metadata /// - **Binary expressions**: field metadata is empty /// - **Boolean expressions**: field metadata is empty - /// - **Cast expressions**: determined by the input expression's field metadata handling + /// - **Cast expressions**: Type-only casts pass through source metadata (stripping extension + /// type keys); casts with explicit target fields use target metadata exactly /// - **Scalar functions**: Generate metadata via function's [`return_field_from_args`] method, /// with the default implementation returning empty field metadata /// - **Aggregate functions**: Generate metadata via function's [`return_field`] method, @@ -602,20 +632,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)), @@ -1290,4 +1316,115 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_and_try_cast_extension_type_metadata() { + use crate::expr::{Cast, TryCast}; + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, + }; + + // Helper to build either Cast or TryCast expression + fn make_cast_expr( + expr: Expr, + target_field: FieldRef, + use_try_cast: bool, + ) -> Expr { + if use_try_cast { + Expr::TryCast(TryCast { + expr: Box::new(expr), + field: target_field, + }) + } else { + Expr::Cast(Cast { + expr: Box::new(expr), + field: target_field, + }) + } + } + + // Run the same test logic for both Cast and TryCast + for use_try_cast in [false, true] { + let cast_name = if use_try_cast { "TryCast" } else { "Cast" }; + + // Create a schema with a field that has extension type metadata + let mut source_meta = HashMap::new(); + source_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ); + source_meta.insert("custom_key".to_string(), "custom_value".to_string()); + + let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta); + + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_field.metadata().clone())); + + // Test 1: Cast to a type without extension metadata strips extension metadata + // but preserves non-extension metadata + let cast_expr = make_cast_expr( + col("foo"), + Arc::new(Field::new("", DataType::Utf8, true)), + use_try_cast, + ); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert!( + result_field + .metadata() + .get(EXTENSION_TYPE_NAME_KEY) + .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" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + + // Test 2: Cast to a field with explicit metadata uses target metadata exactly + let mut target_meta = HashMap::new(); + target_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ); + target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()); + + let target_field = + Field::new("", DataType::Utf8, true).with_metadata(target_meta); + + let cast_expr = + make_cast_expr(col("foo"), Arc::new(target_field), use_try_cast); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()), + "{cast_name}: Extension type name should come from target field" + ); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()), + "{cast_name}: Extension type metadata should come from target field" + ); + assert!( + result_field.metadata().get("custom_key").is_none(), + "{cast_name}: Source metadata should NOT propagate when target has explicit metadata" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + } + } } diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..929a92abcca23 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,9 +161,27 @@ 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)?; + + // 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 let new_expr = if source_type == target_type { - // the argument's data type is already the correct type - source_arg + // 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 { + // Need to create a cast to strip extension metadata + Expr::Cast(datafusion_expr::Cast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + // the argument's data type is already the correct type + source_arg + } } else { // Use an actual cast to get the correct type Expr::Cast(datafusion_expr::Cast { diff --git a/datafusion/functions/src/core/arrow_try_cast.rs b/datafusion/functions/src/core/arrow_try_cast.rs index d27b29ba5736d..0914695b60370 100644 --- a/datafusion/functions/src/core/arrow_try_cast.rs +++ b/datafusion/functions/src/core/arrow_try_cast.rs @@ -26,8 +26,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; @@ -134,8 +134,26 @@ impl ScalarUDFImpl for ArrowTryCastFunc { 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 let new_expr = if source_type == target_type { - source_arg + // 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 { + // Need to create a try_cast to strip extension metadata + Expr::TryCast(datafusion_expr::TryCast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + source_arg + } } else { Expr::TryCast(datafusion_expr::TryCast { expr: Box::new(source_arg), diff --git a/datafusion/physical-expr-adapter/src/rewrite.rs b/datafusion/physical-expr-adapter/src/rewrite.rs index 7345a587ee6a4..3d31feca96031 100644 --- a/datafusion/physical-expr-adapter/src/rewrite.rs +++ b/datafusion/physical-expr-adapter/src/rewrite.rs @@ -319,6 +319,9 @@ mod tests { assert_eq!(source.name(), "__datafusion_file_row_index"); assert_eq!(source.index(), 2); + // The row index column is at index 2, beyond the user-visible schema. + // When the source column lookup fails, the field name is empty. The + // correct field name would be provided by a parent projection/alias. let input_schema = Schema::new(vec![ Field::new("value", DataType::Int64, true), Field::new("__datafusion_file_row_index", DataType::Int64, false) @@ -328,9 +331,11 @@ mod tests { )])), ]); let return_field = expr.return_field(&input_schema)?; - assert_eq!(return_field.name(), "file_row_index"); + // Field name is empty because column index 2 is beyond the schema + assert_eq!(return_field.name(), ""); assert_eq!(return_field.data_type(), &DataType::Int64); assert!(return_field.is_nullable()); + // Exact target field does not preserve source metadata assert!(return_field.metadata().is_empty()); Ok(()) } diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..0588a777230fb 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -51,6 +51,7 @@ proto = [ [dependencies] arrow = { workspace = true } +arrow-schema = { workspace = true } datafusion-common = { workspace = true } datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index dbb91e365af90..d9ee0798433c1 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,8 +23,9 @@ 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 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::{ @@ -59,8 +61,18 @@ fn can_cast_named_struct_types(source: &DataType, target: &DataType) -> bool { pub struct CastExpr { /// The expression to cast pub expr: Arc, - /// Field metadata describing the desired output after casting + /// The target field. + /// + /// For a type-only cast (see [`CastExpr::new`]) this is a field synthesized + /// from the target data type alone and only its data type is meaningful. + /// For a cast built from an explicit field (see + /// [`CastExpr::new_with_target_field`]) its metadata and nullability are + /// applied to the output field as-is. target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata and + /// nullability describe the output field exactly. + explicit_target: bool, /// Cast options cast_options: CastOptions<'static>, } @@ -68,8 +80,12 @@ pub struct CastExpr { // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for CastExpr { fn eq(&self, other: &Self) -> bool { + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. self.expr.eq(&other.expr) - && self.target_field.eq(&other.target_field) + && self.cast_type().eq(other.cast_type()) + && self.target_metadata().eq(&other.target_metadata()) + && self.target_nullable().eq(&other.target_nullable()) && self.cast_options.eq(&other.cast_options) } } @@ -77,7 +93,17 @@ impl PartialEq for CastExpr { impl Hash for CastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.target_field.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } + self.target_nullable().hash(state); self.cast_options.hash(state); } } @@ -85,39 +111,39 @@ impl Hash for CastExpr { impl CastExpr { /// Create a new `CastExpr` using only a `DataType`. /// - /// This constructor is provided for compatibility with existing call sites - /// that only know the target type. It synthesizes a ``Field`` with the - /// given type (**nullable by default**) and no name metadata. Callers that - /// already have a `FieldRef` (for example, coming from schema inference or a - /// resolved column) should prefer [`CastExpr::new_with_target_field`], which - /// preserves the field's name, nullability, and other metadata. In other - /// words: + /// This constructor creates a type-only cast where metadata and nullability + /// are passed through from the source expression (with extension type keys + /// stripped from metadata). This is the most common use case when you only + /// need to change the data type. /// - /// * use `new()` when only a `DataType` is available and you want the legacy - /// semantics of a type-only cast - /// * use `new_with_target_field()` when you need explicit field - /// metadata/name/nullability preserved + /// For explicit control over the output field's metadata and nullability, + /// use [`CastExpr::new_with_target_field`] or the individual builder methods. pub fn new( expr: Arc, cast_type: DataType, cast_options: Option>, ) -> Self { - Self::new_with_target_field( + Self { expr, - cast_type.into_nullable_field_ref(), - cast_options, - ) + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), + } } /// Create a new `CastExpr` with an explicit target `FieldRef`. /// - /// The provided `target_field` is used verbatim for the expression's - /// return schema, so the field's name, nullability, and other metadata are - /// preserved. This is the preferred constructor when the caller already - /// has field information (for example, during logical-to-physical planning). + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// - The field's nullability is preserved + /// + /// This is the preferred constructor when the caller has explicit field + /// information that should be used exactly (for example, during schema + /// enforcement or adapter layers). /// - /// See [`CastExpr::new`] for the compatibility constructor that only accepts - /// a `DataType`. + /// See [`CastExpr::new`] for type-only casts where source metadata should + /// pass through. pub fn new_with_target_field( expr: Arc, target_field: FieldRef, @@ -126,6 +152,7 @@ impl CastExpr { Self { expr, target_field, + explicit_target: true, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), } } @@ -140,7 +167,30 @@ impl CastExpr { self.target_field.data_type() } - /// Field metadata describing the output column after casting. + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// Explicit nullability for the output field, or `None` to pass through source nullability. + pub fn target_nullable(&self) -> Option { + self.explicit_target + .then(|| self.target_field.is_nullable()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. Note that the returned + /// field may not match what `return_field()` returns when evaluated against + /// a schema, since `return_field()` may incorporate source field information. + /// + /// Prefer [`cast_type()`], [`target_metadata()`], and [`target_nullable()`] + /// for direct access to the individual components. + /// + /// [`cast_type()`]: CastExpr::cast_type + /// [`target_metadata()`]: CastExpr::target_metadata + /// [`target_nullable()`]: CastExpr::target_nullable pub fn target_field(&self) -> &FieldRef { &self.target_field } @@ -150,19 +200,53 @@ impl CastExpr { &self.cast_options } + /// Whether this cast has explicit metadata (vs pass-through from source). + pub fn has_explicit_metadata(&self) -> bool { + self.explicit_target + } + + /// Whether this cast has explicit nullability (vs pass-through from source). + pub fn has_explicit_nullability(&self) -> bool { + self.explicit_target + } + 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()), + // Try to get the source field for the name. If the target field is + // explicit, we can fall back to an empty name if the source lookup fails + // (e.g., for virtual row-index columns appended at scan time). + let source_result = self.expr.return_field(input_schema); + + if self.explicit_target { + // Metadata and nullability come from the target field verbatim + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new( + name, + self.cast_type().clone(), + self.target_field.is_nullable(), ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + .with_metadata(self.target_field.metadata().clone()), + )); } + + // Type-only cast: pass through the source metadata and nullability, + // stripping extension type keys (the cast is to a plain storage type). + 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), + ) + }) } /// Check if casting from the specified source type to the target type is a @@ -191,12 +275,6 @@ impl CastExpr { } } -fn is_default_target_field(target_field: &FieldRef) -> bool { - target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty() -} - pub(crate) fn is_order_preserving_cast_family( source_type: &DataType, target_type: &DataType, @@ -267,11 +345,12 @@ impl PhysicalExpr for CastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(CastExpr::new_with_target_field( - Arc::clone(&children[0]), - Arc::clone(&self.target_field), - Some(self.cast_options.clone()), - ))) + Ok(Arc::new(CastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + cast_options: self.cast_options.clone(), + })) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -372,32 +451,61 @@ pub fn cast_with_options( cast_type: DataType, cast_options: Option>, ) -> Result> { - cast_with_target_field( - expr, - input_schema, - cast_type.into_nullable_field_ref(), - cast_options, - ) + let expr_type = expr.data_type(input_schema)?; + + // If the types match, no cast is needed for a type-only cast + if expr_type == cast_type { + return Ok(Arc::clone(&expr)); + } + + let can_build_cast = if requires_nested_struct_cast(&expr_type, &cast_type) { + can_cast_named_struct_types(&expr_type, &cast_type) + } else { + can_cast_types(&expr_type, &cast_type) + }; + + if !can_build_cast { + return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); + } + + Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options))) } /// Return a PhysicalExpression representing `expr` casted to `target_field`, /// 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. +/// If the input expression already has the same data type and the target field +/// has no explicit metadata or nullability constraints, the original expression +/// is returned unchanged. pub fn cast_with_target_field( expr: Arc, input_schema: &Schema, - target_field: FieldRef, + target_field: &FieldRef, cast_options: Option>, ) -> 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)); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata or nullability constraints). This is the field created by + // `into_nullable_field_ref()` when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && 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. + 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 { + return Ok(Arc::clone(&expr)); + } } let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { @@ -415,11 +523,22 @@ pub fn cast_with_target_field( return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); } - Ok(Arc::new(CastExpr::new_with_target_field( - expr, - target_field, - cast_options, - ))) + // For type-only casts, use CastExpr::new which preserves source metadata/nullability. + // For explicit target fields, use new_with_target_field which applies the target's + // extension metadata and nullability. + if is_type_only { + Ok(Arc::new(CastExpr::new( + expr, + cast_type.clone(), + cast_options, + ))) + } else { + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + cast_options, + ))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -990,26 +1109,32 @@ mod tests { #[test] fn field_aware_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); for (child_nullable, target_nullable) in [(true, false), (false, true)] { let schema = Schema::new(vec![Field::new("a", Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", Int64, target_nullable) + .with_metadata(metadata.clone()), + ); let expr = CastExpr::new_with_target_field( col("a", &schema)?, - Arc::new( - Field::new("cast_target", Int64, target_nullable) - .with_metadata(metadata.clone()), - ), + Arc::clone(&target_field), None, ); let field = expr.return_field(&schema)?; - assert_eq!(field.name(), "cast_target"); + // Field name comes from source + assert_eq!(field.name(), "a"); assert_eq!(field.data_type(), &Int64); + // Nullability comes from target assert_eq!(field.is_nullable(), target_nullable); + // Target metadata should be preserved exactly assert_eq!( - field.metadata().get("target_meta").map(String::as_str), - Some("1") + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" ); assert_eq!(expr.nullable(&schema)?, child_nullable || target_nullable); } @@ -1017,6 +1142,38 @@ mod tests { Ok(()) } + #[test] + fn target_field_accessor_returns_the_constructed_field() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", Int32, true)]); + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + let target_field = + Arc::new(Field::new("cast_target", Int64, false).with_metadata(metadata)); + + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + // The field is returned verbatim, including its name. + assert_eq!(expr.target_field(), &target_field); + assert_eq!(expr.cast_type(), &Int64); + assert_eq!(expr.target_metadata(), Some(target_field.metadata())); + assert_eq!(expr.target_nullable(), Some(false)); + assert!(expr.has_explicit_metadata()); + assert!(expr.has_explicit_nullability()); + + // A type-only cast reports no explicit target. + let type_only = CastExpr::new(col("a", &schema)?, Int64, None); + assert_eq!(type_only.cast_type(), &Int64); + assert_eq!(type_only.target_metadata(), None); + assert_eq!(type_only.target_nullable(), None); + assert!(!type_only.has_explicit_metadata()); + assert!(!type_only.has_explicit_nullability()); + + Ok(()) + } + #[test] fn type_only_cast_preserves_legacy_field_name_and_nullability() -> Result<()> { let schema = Schema::new(vec![Field::new("a", Int32, false)]); @@ -1170,7 +1327,8 @@ mod tests { let literal = Arc::new(crate::expressions::Literal::new(ScalarValue::Struct( Arc::new(scalar_struct), ))); - let expr = CastExpr::new_with_target_field(literal, Arc::new(target_field), None); + let target_field = Arc::new(target_field); + let expr = CastExpr::new_with_target_field(literal, target_field, None); let batch = RecordBatch::new_empty(schema); let result = expr.evaluate(&batch)?; @@ -1216,6 +1374,93 @@ mod tests { Ok(()) } + #[test] + fn type_only_cast_strips_extension_metadata() -> Result<()> { + // When using type-only cast (new()), extension metadata from source should NOT propagate + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let expr = CastExpr::new(col("a", &schema)?, Utf8, None); + + let field = expr.return_field(&schema)?; + assert!( + 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" + ); + + Ok(()) + } + + #[test] + fn field_aware_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", Utf8, true).with_metadata(target_meta)); + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware cast should preserve target's non-extension metadata" + ); + + Ok(()) + } + #[test] fn test_check_bigger_cast_precision_loss() { use DataType::*; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 035dd5d5072b0..f5784a10a38e8 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -58,7 +58,7 @@ pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; pub use not::{NotExpr, not}; -pub use try_cast::{TryCastExpr, try_cast}; +pub use try_cast::{TryCastExpr, try_cast, try_cast_with_target_field}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..c054026724fb1 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,40 +23,94 @@ 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 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; use datafusion_common::{Result, not_impl_err}; use datafusion_expr::ColumnarValue; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast -#[derive(Debug, Eq)] +#[derive(Debug, Clone, Eq)] pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + /// The target field. + /// + /// For a type-only cast (see [`TryCastExpr::new`]) this is a field + /// synthesized from the target data type alone and only its data type is + /// meaningful. For a cast built from an explicit field (see + /// [`TryCastExpr::new_with_target_field`]) its metadata is applied to the + /// output field as-is. + target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata + /// describes the output field exactly. + explicit_target: bool, } // 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 + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. + self.expr.eq(&other.expr) + && self.cast_type() == other.cast_type() + && self.target_metadata() == other.target_metadata() } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. + /// + /// This constructor creates a type-only cast where metadata is passed through + /// from the source expression (with extension type keys stripped). + /// TRY_CAST results are always nullable since failed casts return NULL. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self { + expr, + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + } + } + + /// Create a new `TryCastExpr` with an explicit target `FieldRef`. + /// + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// + /// TRY_CAST results are always nullable since failed casts return NULL. + /// + /// See [`TryCastExpr::new`] for type-only casts where source metadata should + /// pass through. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { + expr, + target_field, + explicit_target: true, + } } /// The expression to cast @@ -65,19 +120,33 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target_field.data_type() + } + + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. TRY_CAST results are + /// always nullable regardless of the target field's nullability. + 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 +159,41 @@ 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) + // If metadata is explicit, we can build the field without source + // (though we still try to get source for the name) + let source_result = self.expr.return_field(input_schema); + + if let Some(metadata) = self.target_metadata() { + // Explicit metadata: use it exactly, TRY_CAST is always nullable + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new(name, self.cast_type().clone(), true) + .with_metadata(metadata.clone()), + )); + } + + // Pass-through metadata from source (stripping extension keys) + 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), + ) + }) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +204,17 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( - Arc::clone(&children[0]), - self.cast_type.clone(), - ))) + Ok(Arc::new(TryCastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + })) } 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")] @@ -190,6 +287,60 @@ pub fn try_cast( } } +/// Return a PhysicalExpression representing `expr` casted to `target_field`, +/// preserving any explicit field semantics such as metadata. +/// +/// TRY_CAST results are always nullable since failed casts return NULL. +/// +/// If the input expression already has the same data type, the target field +/// has no explicit metadata constraints, and the source has no extension +/// metadata to strip, the original expression is returned unchanged. +pub fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: &FieldRef, +) -> Result> { + let expr_type = expr.data_type(input_schema)?; + let cast_type = target_field.data_type(); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata constraints). This is the field created by `into_nullable_field_ref()` + // when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && 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. + 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 { + return Ok(Arc::clone(&expr)); + } + } + + if !can_cast_types(&expr_type, cast_type) { + return not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}"); + } + + // For type-only casts, use TryCastExpr::new which preserves source metadata. + // For explicit target fields, use new_with_target_field which applies the target's + // metadata exactly. + if is_type_only { + Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone()))) + } else { + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + ))) + } +} + #[cfg(test)] mod tests { use super::*; @@ -642,6 +793,167 @@ mod tests { Ok(()) } + + #[test] + fn field_aware_try_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", DataType::Utf8, true).with_metadata(target_meta)); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware try_cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware try_cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware try_cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware try_cast should preserve target's non-extension metadata" + ); + // TRY_CAST is always nullable + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). + // TRY_CAST is always nullable regardless of target field's nullability. + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + + for child_nullable in [true, false] { + let schema = + Schema::new(vec![Field::new("a", DataType::Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", DataType::Int64, false) // target says non-nullable + .with_metadata(metadata.clone()), + ); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + // Field name comes from source + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + // TRY_CAST is ALWAYS nullable (ignores target field's nullability) + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + // Target metadata should be preserved exactly + assert_eq!( + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" + ); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + } + + Ok(()) + } + + #[test] + fn type_only_try_cast_strips_extension_keys() -> Result<()> { + // Type-only cast should strip extension keys but preserve other source metadata + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.extension".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "ext_meta".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false).with_metadata(source_meta), + ]); + + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + let field = expr.return_field(&schema)?; + + // Extension keys should be stripped + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only try_cast should strip extension type name" + ); + assert!( + 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" + ); + // Field name preserved, type changed, always nullable + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn type_only_try_cast_is_always_nullable() -> Result<()> { + // TRY_CAST is always nullable even when source is non-nullable + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + + Ok(()) + } } #[cfg(all(test, feature = "proto"))] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index f80d1b15bdc59..91c36cb09cd35 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, @@ -41,7 +41,7 @@ use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::var_provider::VarType; use datafusion_expr::var_provider::is_system_variables; use datafusion_expr::{ - Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, TryCast, binary_expr, lit, + Between, BinaryExpr, Expr, Like, Operator, TryCast, binary_expr, lit, }; /// [PhysicalExpr] evaluate DataFusion expressions such as `A + 1`, or `CAST(c1 @@ -387,23 +387,11 @@ pub fn create_physical_expr( Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, - Arc::clone(field), + field, 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(), + field, ) } Expr::Not(expr) => expressions::not(create_physical_expr( @@ -744,6 +732,7 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { mod tests { use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field}; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_expr::col; use super::*; @@ -768,6 +757,14 @@ mod tests { .expect("planner should lower logical CAST to CastExpr") } + fn as_planner_try_cast( + physical: &Arc, + ) -> &expressions::TryCastExpr { + physical + .downcast_ref::() + .expect("planner should lower logical TRY_CAST to TryCastExpr") + } + #[test] fn test_create_physical_expr_scalar_input_output() -> Result<()> { let expr = col("letter").eq(lit("A")); @@ -801,9 +798,21 @@ mod tests { #[test] fn test_cast_lowering_preserves_target_field_metadata() -> Result<()> { let schema = test_cast_schema(); + + // Target field with both extension metadata and custom metadata. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( - Field::new("cast_target", DataType::Int64, true) - .with_metadata([("target_meta".to_string(), "1".to_string())].into()), + Field::new("cast_target", DataType::Int64, true).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ), + (EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()), + ("custom_target_meta".to_string(), "custom_value".to_string()), + ] + .into(), + ), ); let cast_expr = Expr::Cast(Cast::new_from_field( Box::new(col("a")), @@ -813,8 +822,35 @@ mod tests { let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); + // The CastExpr stores the target type and all target metadata + assert_eq!(cast.cast_type(), &DataType::Int64); + let target_metadata = cast.target_metadata().expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + assert_eq!(cast.target_nullable(), Some(true)); + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_target_meta"), + Some(&"custom_value".to_string()), + "All target metadata should propagate with exact semantics" + ); assert!(physical.nullable(&schema)?); Ok(()) @@ -840,22 +876,85 @@ mod tests { #[test] fn test_cast_lowering_preserves_same_type_field_semantics() -> Result<()> { let schema = test_cast_schema(); + + // Same-type cast with extension metadata on target. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( Field::new("same_type_cast", DataType::Int32, true).with_metadata( - [("target_meta".to_string(), "same-type".to_string())].into(), + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.opaque".to_string(), + ), + ("custom_meta".to_string(), "custom_value".to_string()), + ] + .into(), ), ); - let cast_expr = Expr::Cast(Cast::new_from_field( - Box::new(col("a")), - Arc::clone(&target_field), - )); - let physical = lower_cast_expr(&cast_expr, &schema)?; - let cast = as_planner_cast(&physical); + for use_try_cast in [false, true] { + // For error labelling + let cast_name = if use_try_cast { "TRY_CAST" } else { "CAST" }; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + let cast_expr = if use_try_cast { + Expr::TryCast(TryCast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + } else { + Expr::Cast(Cast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + }; + + let physical = lower_cast_expr(&cast_expr, &schema)?; + + // Extract common fields - both CastExpr and TryCastExpr have these + let (cast_type, target_metadata, target_nullable) = if use_try_cast { + let cast = as_planner_try_cast(&physical); + (cast.cast_type(), cast.target_metadata(), None) + } else { + let cast = as_planner_cast(&physical); + ( + cast.cast_type(), + cast.target_metadata(), + cast.target_nullable(), + ) + }; + + // Verify the physical expression stores correct metadata (same for both) + assert_eq!(cast_type, &DataType::Int32, "{cast_name}: cast_type"); + let target_metadata = target_metadata.expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: extension type name" + ); + + // Only CastExpr tracks target_nullable (TryCast is always nullable) + if !use_try_cast { + assert_eq!(target_nullable, Some(true), "{cast_name}: target_nullable"); + } + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: return_field extension type name" + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_meta"), + Some(&"custom_value".to_string()), + "{cast_name}: All target metadata should propagate with exact semantics" + ); + assert!( + physical.nullable(&schema)?, + "{cast_name}: should be nullable" + ); + } Ok(()) } diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index d4d0ea7292ffe..2830edd35760f 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -429,7 +429,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field)))) } ExprType::TryCast(cast) => { @@ -442,7 +443,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::TryCast(TryCast::new_from_field( expr, Arc::new(field), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index a450f7a7e888f..4dc3291f31d57 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -2693,6 +2693,19 @@ fn roundtrip_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::Cast(Cast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); } #[test] @@ -2703,6 +2716,19 @@ fn roundtrip_try_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::TryCast(TryCast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); + let test_expr = Expr::TryCast(TryCast::new(Box::new(lit("not a bool")), DataType::Boolean)); diff --git a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..01a19454e9a80 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -45,5 +45,55 @@ 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 extension type should also preserve extension metadata +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 + +# TRY_CAST to UUID from a subquery +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 + +# arrow_cast from UUID to same underlying type (FixedSizeBinary(16)) strips +# extension metadata (type-only cast semantics) +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 + +# arrow_cast to a different type strips extension metadata (type-only cast semantics) +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