Allow logical expressions to express a cast to an extension type

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -712,10 +712,7 @@ impl ScalarUDFImpl for CastToI64UDF {
             arg
         } else {
             // need to use an actual cast to get the correct type
-            Expr::Cast(datafusion_expr::Cast {
-                expr: Box::new(arg),
-                data_type: DataType::Int64,
-            })
+            Expr::Cast(datafusion_expr::Cast::new(Box::new(arg), DataType::Int64))
         };
         // return the newly written argument to DataFusion
         Ok(ExprSimplifyResult::Simplified(new_expr))

datafusion/expr/src/expr.rs

@@ -1012,13 +1012,23 @@ pub struct Cast {
     /// The expression being cast
     pub expr: Box<Expr>,
     /// The `DataType` the expression will yield
-    pub data_type: DataType,
+    pub data_type: FieldRef,
 }
 
 impl Cast {
     /// Create a new Cast expression
     pub fn new(expr: Box<Expr>, data_type: DataType) -> Self {
-        Self { expr, data_type }
+        Self {
+            expr,
+            data_type: Field::new("", data_type, true).into(),
+        }
+    }
+
+    pub fn new_from_field(expr: Box<Expr>, field: FieldRef) -> Self {
+        Self {
+            expr,
+            data_type: field,
+        }
     }
 }
 
@@ -1028,13 +1038,23 @@ pub struct TryCast {
     /// The expression being cast
     pub expr: Box<Expr>,
     /// The `DataType` the expression will yield
-    pub data_type: DataType,
+    pub data_type: FieldRef,
 }
 
 impl TryCast {
     /// Create a new TryCast expression
     pub fn new(expr: Box<Expr>, data_type: DataType) -> Self {
-        Self { expr, data_type }
+        Self {
+            expr,
+            data_type: Field::new("", data_type, true).into(),
+        }
+    }
+
+    pub fn new_from_field(expr: Box<Expr>, field: FieldRef) -> Self {
+        Self {
+            expr,
+            data_type: field,
+        }
     }
 }
 
@@ -3488,10 +3508,28 @@ impl Display for Expr {
                 write!(f, "END")
             }
             Expr::Cast(Cast { expr, data_type }) => {
-                write!(f, "CAST({expr} AS {data_type})")
+                if data_type.metadata().is_empty() {
+                    write!(f, "CAST({expr} AS {})", data_type.data_type())
+                } else {
+                    write!(
+                        f,
+                        "CAST({expr} AS {}<{:?}>)",
+                        data_type.data_type(),
+                        data_type.metadata()
+                    )
+                }
             }
             Expr::TryCast(TryCast { expr, data_type }) => {
-                write!(f, "TRY_CAST({expr} AS {data_type})")
+                if data_type.metadata().is_empty() {
+                    write!(f, "TRY_CAST({expr} AS {})", data_type.data_type())
+                } else {
+                    write!(
+                        f,
+                        "TRY_CAST({expr} AS {}<{:?}>)",
+                        data_type.data_type(),
+                        data_type.metadata()
+                    )
+                }
             }
             Expr::Not(expr) => write!(f, "NOT {expr}"),
             Expr::Negative(expr) => write!(f, "(- {expr})"),
@@ -3844,7 +3882,7 @@ mod test {
     fn format_cast() -> Result<()> {
         let expr = Expr::Cast(Cast {
             expr: Box::new(Expr::Literal(ScalarValue::Float32(Some(1.23)), None)),
-            data_type: DataType::Utf8,
+            data_type: Field::new("", DataType::Utf8, true).into(),
         });
         let expected_canonical = "CAST(Float32(1.23) AS Utf8)";
         assert_eq!(expected_canonical, format!("{expr}"));

datafusion/expr/src/expr_schema.rs

@@ -127,7 +127,9 @@ impl ExprSchemable for Expr {
                     .map_or(Ok(DataType::Null), |e| e.get_type(schema))
             }
             Expr::Cast(Cast { data_type, .. })
-            | Expr::TryCast(TryCast { data_type, .. }) => Ok(data_type.clone()),
+            | Expr::TryCast(TryCast { data_type, .. }) => {
+                Ok(data_type.data_type().clone())
+            }
             Expr::Unnest(Unnest { expr }) => {
                 let arg_data_type = expr.get_type(schema)?;
                 // Unnest's output type is the inner type of the list
@@ -592,7 +594,14 @@ impl ExprSchemable for Expr {
             // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)),
             Expr::Cast(Cast { expr, data_type }) => expr
                 .to_field(schema)
-                .map(|(_, f)| f.as_ref().clone().with_data_type(data_type.clone()))
+                .map(|(_, f)| {
+                    f.as_ref()
+                        .clone()
+                        .with_data_type(data_type.data_type().clone())
+                        .with_metadata(f.metadata().clone())
+                    // TODO: should nullability be overridden here or derived from the
+                    // input expression?
+                })
                 .map(Arc::new),
             Expr::Like(_)
             | Expr::SimilarTo(_)

datafusion/expr/src/tree_node.rs

@@ -222,10 +222,10 @@ impl TreeNode for Expr {
                 }),
             Expr::Cast(Cast { expr, data_type }) => expr
                 .map_elements(f)?
-                .update_data(|be| Expr::Cast(Cast::new(be, data_type))),
+                .update_data(|be| Expr::Cast(Cast::new_from_field(be, data_type))),
             Expr::TryCast(TryCast { expr, data_type }) => expr
                 .map_elements(f)?
-                .update_data(|be| Expr::TryCast(TryCast::new(be, data_type))),
+                .update_data(|be| Expr::TryCast(TryCast::new_from_field(be, data_type))),
             Expr::ScalarFunction(ScalarFunction { func, args }) => {
                 args.map_elements(f)?.map_data(|new_args| {
                     Ok(Expr::ScalarFunction(ScalarFunction::new_udf(

datafusion/functions/src/core/arrow_cast.rs

@@ -161,7 +161,7 @@ impl ScalarUDFImpl for ArrowCastFunc {
             // Use an actual cast to get the correct type
             Expr::Cast(datafusion_expr::Cast {
                 expr: Box::new(arg),
-                data_type: target_type,
+                data_type: Field::new("", target_type, true).into(),
             })
         };
         // return the newly written argument to DataFusion

datafusion/physical-expr/src/planner.rs

@@ -291,12 +291,14 @@ pub fn create_physical_expr(
         Expr::Cast(Cast { expr, data_type }) => expressions::cast(
             create_physical_expr(expr, input_dfschema, execution_props)?,
             input_schema,
-            data_type.clone(),
+            // TODO: this drops extension metadata associated with the cast
+            data_type.data_type().clone(),
         ),
         Expr::TryCast(TryCast { expr, data_type }) => expressions::try_cast(
             create_physical_expr(expr, input_dfschema, execution_props)?,
             input_schema,
-            data_type.clone(),
+            // TODO: this drops extension metadata associated with the cast
+            data_type.data_type().clone(),
         ),
         Expr::Not(expr) => {
             expressions::not(create_physical_expr(expr, input_dfschema, execution_props)?)

datafusion/proto/proto/datafusion.proto

@@ -586,11 +586,15 @@ message WhenThen {
 message CastNode {
   LogicalExprNode expr = 1;
   datafusion_common.ArrowType arrow_type = 2;
+  map<string, string> metadata = 3;
+  optional bool nullable = 4;
 }
 
 message TryCastNode {
   LogicalExprNode expr = 1;
   datafusion_common.ArrowType arrow_type = 2;
+  map<string, string> metadata = 3;
+  optional bool nullable = 4;
 }
 
 message SortExprNode {