Add Extension Type / Metadata support for Scalar UDFs

datafusion/common/src/dfschema.rs

@@ -472,7 +472,7 @@ impl DFSchema {
         let matches = self.qualified_fields_with_unqualified_name(name);
         match matches.len() {
             0 => Err(unqualified_field_not_found(name, self)),
-            1 => Ok((matches[0].0, (matches[0].1))),
+            1 => Ok((matches[0].0, matches[0].1)),
             _ => {
                 // When `matches` size > 1, it doesn't necessarily mean an `ambiguous name` problem.
                 // Because name may generate from Alias/... . It means that it don't own qualifier.
@@ -969,16 +969,28 @@ impl Display for DFSchema {
 /// widely used in the DataFusion codebase.
 pub trait ExprSchema: std::fmt::Debug {
     /// Is this column reference nullable?
-    fn nullable(&self, col: &Column) -> Result<bool>;
+    fn nullable(&self, col: &Column) -> Result<bool> {
+        Ok(self.to_field(col)?.is_nullable())
+    }
 
     /// What is the datatype of this column?
-    fn data_type(&self, col: &Column) -> Result<&DataType>;
+    fn data_type(&self, col: &Column) -> Result<&DataType> {
+        Ok(self.to_field(col)?.data_type())
+    }
 
     /// Returns the column's optional metadata.
-    fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>>;
+    fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
+        Ok(self.to_field(col)?.metadata())
+    }
 
     /// Return the column's datatype and nullability
-    fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)>;
+    fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)> {
+        let field = self.to_field(col)?;
+        Ok((field.data_type(), field.is_nullable()))
+    }
+
+    // Return the column's field
+    fn to_field(&self, col: &Column) -> Result<&Field>;
 }
 
 // Implement `ExprSchema` for `Arc<DFSchema>`
@@ -998,24 +1010,15 @@ impl<P: AsRef<DFSchema> + std::fmt::Debug> ExprSchema for P {
     fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)> {
         self.as_ref().data_type_and_nullable(col)
     }
-}
 
-impl ExprSchema for DFSchema {
-    fn nullable(&self, col: &Column) -> Result<bool> {
-        Ok(self.field_from_column(col)?.is_nullable())
-    }
-
-    fn data_type(&self, col: &Column) -> Result<&DataType> {
-        Ok(self.field_from_column(col)?.data_type())
-    }
-
-    fn metadata(&self, col: &Column) -> Result<&HashMap<String, String>> {
-        Ok(self.field_from_column(col)?.metadata())
+    fn to_field(&self, col: &Column) -> Result<&Field> {
+        self.as_ref().to_field(col)
     }
+}
 
-    fn data_type_and_nullable(&self, col: &Column) -> Result<(&DataType, bool)> {
-        let field = self.field_from_column(col)?;
-        Ok((field.data_type(), field.is_nullable()))
+impl ExprSchema for DFSchema {
+    fn to_field(&self, col: &Column) -> Result<&Field> {
+        self.field_from_column(col)
     }
 }
 

datafusion/core/Cargo.toml

@@ -95,7 +95,7 @@ extended_tests = []
 [dependencies]
 arrow = { workspace = true }
 arrow-ipc = { workspace = true }
-arrow-schema = { workspace = true }
+arrow-schema = { workspace = true, features = ["canonical_extension_types"] }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 bzip2 = { version = "0.5.2", optional = true }

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -128,7 +128,7 @@ fn test_update_matching_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            DataType::Int32,
+            Field::new("f", DataType::Int32, true),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 2))),
@@ -193,7 +193,7 @@ fn test_update_matching_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            DataType::Int32,
+            Field::new("f", DataType::Int32, true),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 3))),
@@ -261,7 +261,7 @@ fn test_update_projected_exprs() -> Result<()> {
                     Arc::new(Column::new("b", 1)),
                 )),
             ],
-            DataType::Int32,
+            Field::new("f", DataType::Int32, true),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d", 2))),
@@ -326,7 +326,7 @@ fn test_update_projected_exprs() -> Result<()> {
                     Arc::new(Column::new("b_new", 1)),
                 )),
             ],
-            DataType::Int32,
+            Field::new("f", DataType::Int32, true),
         )),
         Arc::new(CaseExpr::try_new(
             Some(Arc::new(Column::new("d_new", 3))),

datafusion/core/tests/tpc-ds/49.sql

@@ -110,7 +110,7 @@ select  channel, item, return_ratio, return_rank, currency_rank from
  		where 
  			sr.sr_return_amt > 10000 
  			and sts.ss_net_profit > 1
-                         and sts.ss_net_paid > 0 
+                         and sts.ss_net_paid > 0
                          and sts.ss_quantity > 0
                          and ss_sold_date_sk = d_date_sk
                          and d_year = 2000