Skip to content

fix: make a cast target's metadata authoritative #24864

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
adriangb wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: make a cast target's metadata authoritative #24864

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 42 additions & 7 deletions datafusion/expr/src/expr_rewriter/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
Expand All @@ -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),
}
}
}
Expand All @@ -283,6 +288,36 @@ fn coerce_exprs_for_schema(
.collect::<Result<_>>()
}

/// 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<Expr> {
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 {
Expand Down
34 changes: 10 additions & 24 deletions datafusion/expr/src/expr_schema.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -91,27 +90,14 @@ fn cast_output_field(
target_field: &FieldRef,
force_nullable: bool,
) -> Arc<Field> {
// 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);
}
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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!(
Expand Down
15 changes: 6 additions & 9 deletions datafusion/functions/src/core/arrow_cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
33 changes: 13 additions & 20 deletions datafusion/physical-expr/src/expressions/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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()),
)
})
}
Expand Down Expand Up @@ -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));
}
}
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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(())
Expand Down
32 changes: 12 additions & 20 deletions datafusion/physical-expr/src/expressions/try_cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()),
)
})
}
Expand Down Expand Up @@ -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));
}
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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");
Expand Down
Loading