feat: enhance TryCastExpr with target field semantics

- Added a target_field attribute to TryCastExpr to store field metadata for the desired output.
- Introduced methods for creating TryCastExpr instances with and without explicit target field information, ensuring that casting semantics are preserved.
- Updated the return_field and cast_type methods to reflect changes in field metadata handling.
- Enhanced the try_cast and try_cast_with_target_field functions to accept and utilize the target field.
- Added tests to validate the preservation of target field semantics in TryCastExpr operations and enhanced integration with the logical to physical expression conversion.

Author: kosiew

datafusion/physical-expr/src/expressions/mod.rs

@@ -55,5 +55,6 @@ pub use literal::{Literal, lit};
 pub use negative::{NegativeExpr, negative};
 pub use no_op::NoOp;
 pub use not::{NotExpr, not};
+pub(crate) use try_cast::try_cast_with_target_field;
 pub use try_cast::{TryCastExpr, try_cast};
 pub use unknown_column::UnKnownColumn;

datafusion/physical-expr/src/expressions/try_cast.rs

@@ -26,6 +26,7 @@ use arrow::compute::CastOptions;
 use arrow::datatypes::{DataType, FieldRef, Schema};
 use arrow::record_batch::RecordBatch;
 use compute::can_cast_types;
+use datafusion_common::datatype::DataTypeExt;
 use datafusion_common::format::DEFAULT_FORMAT_OPTIONS;
 use datafusion_common::{Result, not_impl_err};
 use datafusion_expr::ColumnarValue;
@@ -35,28 +36,36 @@ use datafusion_expr::ColumnarValue;
 pub struct TryCastExpr {
     /// The expression to cast
     expr: Arc<dyn PhysicalExpr>,
-    /// The data type to cast to
-    cast_type: DataType,
+    /// Field metadata describing the desired output after casting
+    target_field: FieldRef,
 }
 
 // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808
 impl PartialEq for TryCastExpr {
     fn eq(&self, other: &Self) -> bool {
-        self.expr.eq(&other.expr) && self.cast_type == other.cast_type
+        self.expr.eq(&other.expr) && self.target_field == other.target_field
     }
 }
 
 impl Hash for TryCastExpr {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.expr.hash(state);
-        self.cast_type.hash(state);
+        self.target_field.hash(state);
     }
 }
 
 impl TryCastExpr {
     /// Create a new CastExpr
     pub fn new(expr: Arc<dyn PhysicalExpr>, cast_type: DataType) -> Self {
-        Self { expr, cast_type }
+        Self::new_with_target_field(expr, cast_type.into_nullable_field_ref())
+    }
+
+    /// Create a new TryCastExpr with an explicit target field.
+    pub fn new_with_target_field(
+        expr: Arc<dyn PhysicalExpr>,
+        target_field: FieldRef,
+    ) -> Self {
+        Self { expr, target_field }
     }
 
     /// The expression to cast
@@ -66,13 +75,45 @@ impl TryCastExpr {
 
     /// The data type to cast to
     pub fn cast_type(&self) -> &DataType {
-        &self.cast_type
+        self.target_field.data_type()
+    }
+
+    /// Field metadata describing the output column after casting.
+    pub fn target_field(&self) -> &FieldRef {
+        &self.target_field
+    }
+
+    fn is_default_target_field(&self) -> bool {
+        self.target_field.name().is_empty()
+            && self.target_field.is_nullable()
+            && self.target_field.metadata().is_empty()
+    }
+
+    fn resolved_target_field(&self, input_schema: &Schema) -> Result<FieldRef> {
+        if self.is_default_target_field() {
+            self.expr.return_field(input_schema).map(|field| {
+                Arc::new(
+                    field
+                        .as_ref()
+                        .clone()
+                        .with_data_type(self.cast_type().clone())
+                        .with_nullable(true),
+                )
+            })
+        } else {
+            Ok(Arc::clone(&self.target_field))
+        }
+    }
+
+    fn preserves_child_field_semantics(&self, input_schema: &Schema) -> Result<bool> {
+        Ok(self.resolved_target_field(input_schema)?
+            == self.expr.return_field(input_schema)?)
     }
 }
 
 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())
     }
 }
 
@@ -83,7 +124,7 @@ impl PhysicalExpr for TryCastExpr {
     }
 
     fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
-        Ok(self.cast_type.clone())
+        Ok(self.cast_type().clone())
     }
 
     fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
@@ -96,14 +137,11 @@ 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<FieldRef> {
-        self.expr
-            .return_field(input_schema)
-            .map(|f| f.as_ref().clone().with_data_type(self.cast_type.clone()))
-            .map(Arc::new)
+        self.resolved_target_field(input_schema)
     }
 
     fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
@@ -114,16 +152,16 @@ impl PhysicalExpr for TryCastExpr {
         self: Arc<Self>,
         children: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        Ok(Arc::new(TryCastExpr::new(
+        Ok(Arc::new(TryCastExpr::new_with_target_field(
             Arc::clone(&children[0]),
-            self.cast_type.clone(),
+            Arc::clone(&self.target_field),
         )))
     }
 
     fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "TRY_CAST(")?;
         self.expr.fmt_sql(f)?;
-        write!(f, " AS {:?})", self.cast_type)
+        write!(f, " AS {:?})", self.cast_type())
     }
 }
 
@@ -135,12 +173,26 @@ pub fn try_cast(
     expr: Arc<dyn PhysicalExpr>,
     input_schema: &Schema,
     cast_type: DataType,
+) -> Result<Arc<dyn PhysicalExpr>> {
+    try_cast_with_target_field(expr, input_schema, cast_type.into_nullable_field_ref())
+}
+
+pub(crate) fn try_cast_with_target_field(
+    expr: Arc<dyn PhysicalExpr>,
+    input_schema: &Schema,
+    target_field: FieldRef,
 ) -> Result<Arc<dyn PhysicalExpr>> {
     let expr_type = expr.data_type(input_schema)?;
-    if expr_type == cast_type {
-        Ok(Arc::clone(&expr))
-    } else if can_cast_types(&expr_type, &cast_type) {
-        Ok(Arc::new(TryCastExpr::new(expr, cast_type)))
+    let cast_type = target_field.data_type().clone();
+    let candidate = TryCastExpr::new_with_target_field(Arc::clone(&expr), target_field);
+
+    if expr_type == cast_type
+        && (candidate.is_default_target_field()
+            || candidate.preserves_child_field_semantics(input_schema)?)
+    {
+        Ok(expr)
+    } else if expr_type == cast_type || can_cast_types(&expr_type, &cast_type) {
+        Ok(Arc::new(candidate))
     } else {
         not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}")
     }
@@ -161,6 +213,7 @@ mod tests {
         datatypes::*,
     };
     use datafusion_physical_expr_common::physical_expr::fmt_sql;
+    use std::collections::HashMap;
 
     // runs an end-to-end test of physical type cast
     // 1. construct a record batch with a column "a" of type A
@@ -564,6 +617,51 @@ mod tests {
         result.expect_err("expected Invalid TRY_CAST");
     }
 
+    #[test]
+    fn field_aware_same_type_try_cast_preserves_explicit_target_field() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+        let a_col = col("a", &schema)?;
+        let logical_field = Arc::new(
+            Field::new("logical_a", DataType::Int32, true).with_metadata(HashMap::from(
+                [("target_meta".to_string(), "1".to_string())],
+            )),
+        );
+        let expr = try_cast_with_target_field(a_col, &schema, logical_field)?;
+
+        let try_cast_expr = expr
+            .as_any()
+            .downcast_ref::<TryCastExpr>()
+            .expect("explicit same-type target should preserve TryCastExpr");
+        let field = try_cast_expr.return_field(&schema)?;
+
+        assert_eq!(field.name(), "logical_a");
+        assert!(field.is_nullable());
+        assert_eq!(
+            field.metadata().get("target_meta").map(String::as_str),
+            Some("1")
+        );
+        assert!(expr.nullable(&schema)?);
+
+        Ok(())
+    }
+
+    #[test]
+    fn default_same_type_try_cast_is_elided() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+        let a_col = col("a", &schema)?;
+        let target_field = DataType::Int32.into_nullable_field_ref();
+        let expr = try_cast_with_target_field(a_col, &schema, target_field)?;
+
+        assert!(
+            expr.as_any()
+                .downcast_ref::<crate::expressions::Column>()
+                .is_some()
+        );
+        assert!(expr.as_any().downcast_ref::<TryCastExpr>().is_none());
+
+        Ok(())
+    }
+
     // create decimal array with the specified precision and scale
     fn create_decimal_array(array: &[i128], precision: u8, scale: i8) -> Decimal128Array {
         let mut decimal_builder = Decimal128Builder::with_capacity(array.len());

datafusion/physical-expr/src/planner.rs

@@ -25,7 +25,7 @@ use crate::{
 
 use arrow::datatypes::Schema;
 use datafusion_common::config::ConfigOptions;
-use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata};
+use datafusion_common::metadata::FieldMetadata;
 use datafusion_common::{
     DFSchema, Result, ScalarValue, ToDFSchema, exec_err, not_impl_err, plan_err,
 };
@@ -34,7 +34,7 @@ use datafusion_expr::expr::{Alias, Cast, InList, Placeholder, ScalarFunction};
 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
@@ -298,22 +298,10 @@ pub fn create_physical_expr(
             )
         }
         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, execution_props)?,
                 input_schema,
-                field.data_type().clone(),
+                Arc::clone(field),
             )
         }
         Expr::Not(expr) => {
@@ -433,7 +421,7 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc<dyn PhysicalExpr> {
 
 #[cfg(test)]
 mod tests {
-    use crate::expressions::{CastExpr, Column};
+    use crate::expressions::{CastExpr, Column, TryCastExpr};
     use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray};
     use arrow::datatypes::{DataType, Field};
     use datafusion_common::datatype::DataTypeExt;
@@ -498,10 +486,12 @@ mod tests {
             Box::new(expr.clone()),
             Arc::clone(&extension_field_type),
         ));
-        let err =
-            create_physical_expr(&try_cast_expr, &df_schema, &ExecutionProps::new())
-                .unwrap_err();
-        assert!(err.message().contains(extension_name));
+        let physical_expr =
+            create_physical_expr(&try_cast_expr, &df_schema, &ExecutionProps::new())?;
+        let field = physical_expr.return_field(&schema)?;
+        assert_eq!(field.data_type(), extension_field_type.data_type());
+        assert_eq!(field.metadata(), extension_field_type.metadata());
+        assert!(physical_expr.nullable(&schema)?);
 
         Ok(())
     }
@@ -607,6 +597,43 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_create_physical_expr_try_cast_preserves_target_field_semantics() -> Result<()>
+    {
+        let schema = Schema::new(vec![Field::new("value", DataType::Utf8, false)]);
+        let df_schema = DFSchema::try_from(schema.clone())?;
+        let target_field = Arc::new(
+            Field::new("logical_try_cast", DataType::Int64, true).with_metadata(
+                HashMap::from([("semantic".to_string(), "target".to_string())]),
+            ),
+        );
+        let expr = Expr::TryCast(TryCast::new_from_field(
+            Box::new(col("value")),
+            Arc::clone(&target_field),
+        ));
+
+        let physical_expr =
+            create_physical_expr(&expr, &df_schema, &ExecutionProps::new())?;
+        let try_cast_expr = physical_expr
+            .as_any()
+            .downcast_ref::<TryCastExpr>()
+            .expect("logical try_cast should lower to TryCastExpr");
+
+        assert_eq!(try_cast_expr.target_field(), &target_field);
+
+        let field = physical_expr.return_field(&schema)?;
+        assert_eq!(field.name(), "logical_try_cast");
+        assert_eq!(field.data_type(), &DataType::Int64);
+        assert!(field.is_nullable());
+        assert_eq!(
+            field.metadata().get("semantic").map(String::as_str),
+            Some("target")
+        );
+        assert!(physical_expr.nullable(&schema)?);
+
+        Ok(())
+    }
+
     /// Test that deeply nested expressions do not cause a stack overflow.
     ///
     /// This test only runs when the `recursive_protection` feature is enabled,