fix: propagate field metadata through coalesce, nvl, and nvl2

datafusion/functions/src/core/coalesce.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::utils::unanimous_metadata;
 use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::{Result, exec_err, internal_err, plan_err};
 use datafusion_expr::binary::try_type_union_resolution;
@@ -86,7 +87,12 @@ impl ScalarUDFImpl for CoalesceFunc {
             .find_or_first(|d| !d.is_null())
             .unwrap()
             .clone();
-        Ok(Field::new(self.name(), return_type, nullable).into())
+        // Propagate field metadata (e.g. Arrow extension types) only when the
+        // arguments that can supply a value agree on it.
+        let metadata = unanimous_metadata(args.arg_fields.iter().map(|f| f.as_ref()));
+        Ok(Field::new(self.name(), return_type, nullable)
+            .with_metadata(metadata)
+            .into())
     }
 
     fn simplify(
@@ -146,3 +152,97 @@ impl ScalarUDFImpl for CoalesceFunc {
         self.doc()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashMap;
+    use std::sync::Arc;
+
+    const EXTENSION_KEY: &str = "ARROW:extension:name";
+
+    fn field(name: &str, dt: DataType, nullable: bool) -> FieldRef {
+        Arc::new(Field::new(name, dt, nullable))
+    }
+
+    fn ext_field(name: &str, dt: DataType, extension_name: &str) -> FieldRef {
+        Arc::new(Field::new(name, dt, true).with_metadata(HashMap::from([(
+            EXTENSION_KEY.to_string(),
+            extension_name.to_string(),
+        )])))
+    }
+
+    fn return_field(arg_fields: &[FieldRef]) -> FieldRef {
+        let scalars = vec![None; arg_fields.len()];
+        CoalesceFunc::new()
+            .return_field_from_args(ReturnFieldArgs {
+                arg_fields,
+                scalar_arguments: &scalars,
+            })
+            .unwrap()
+    }
+
+    #[test]
+    fn propagates_metadata_when_arguments_agree() {
+        let ret = return_field(&[
+            ext_field("a", DataType::Binary, "geoarrow.wkb"),
+            ext_field("b", DataType::Binary, "geoarrow.wkb"),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Binary);
+        assert!(ret.is_nullable());
+        assert_eq!(
+            ret.metadata().get(EXTENSION_KEY).map(String::as_str),
+            Some("geoarrow.wkb")
+        );
+    }
+
+    #[test]
+    fn null_literal_argument_does_not_block_metadata() {
+        // An untyped NULL argument carries no metadata but cannot contribute
+        // a typed value either, so it does not participate in the agreement.
+        let ret = return_field(&[
+            field("null", DataType::Null, true),
+            ext_field("b", DataType::Binary, "geoarrow.wkb"),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Binary);
+        assert_eq!(
+            ret.metadata().get(EXTENSION_KEY).map(String::as_str),
+            Some("geoarrow.wkb")
+        );
+    }
+
+    #[test]
+    fn drops_metadata_when_arguments_disagree() {
+        let ret = return_field(&[
+            ext_field("a", DataType::Binary, "geoarrow.wkb"),
+            field("plain", DataType::Binary, true),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Binary);
+        assert!(ret.metadata().is_empty());
+    }
+
+    #[test]
+    fn no_metadata_when_all_arguments_are_null() {
+        let ret = return_field(&[
+            field("a", DataType::Null, true),
+            field("b", DataType::Null, true),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Null);
+        assert!(ret.metadata().is_empty());
+    }
+
+    #[test]
+    fn nullability_is_unchanged_by_metadata_propagation() {
+        let ret = return_field(&[
+            ext_field("a", DataType::Binary, "geoarrow.wkb"),
+            Arc::new(Field::new("b", DataType::Binary, false).with_metadata(
+                HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]),
+            )),
+        ]);
+        assert!(!ret.is_nullable());
+        assert_eq!(
+            ret.metadata().get(EXTENSION_KEY).map(String::as_str),
+            Some("geoarrow.wkb")
+        );
+    }
+}

datafusion/functions/src/core/nvl2.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::utils::unanimous_metadata;
 use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::{Result, internal_err, utils::take_function_args};
 use datafusion_expr::{
@@ -94,7 +95,15 @@ impl ScalarUDFImpl for NVL2Func {
         let nullable =
             args.arg_fields[1].is_nullable() || args.arg_fields[2].is_nullable();
         let return_type = args.arg_fields[1].data_type().clone();
-        Ok(Field::new(self.name(), return_type, nullable).into())
+        // The result value comes from the second or third argument; propagate
+        // field metadata (e.g. Arrow extension types) only when they agree on
+        // it. The first argument is only tested for NULL and does not
+        // contribute a value.
+        let metadata =
+            unanimous_metadata(args.arg_fields[1..3].iter().map(|f| f.as_ref()));
+        Ok(Field::new(self.name(), return_type, nullable)
+            .with_metadata(metadata)
+            .into())
     }
 
     fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -145,3 +154,56 @@ impl ScalarUDFImpl for NVL2Func {
         self.doc()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashMap;
+    use std::sync::Arc;
+
+    const EXTENSION_KEY: &str = "ARROW:extension:name";
+
+    fn ext_field(name: &str, dt: DataType, extension_name: &str) -> FieldRef {
+        Arc::new(Field::new(name, dt, true).with_metadata(HashMap::from([(
+            EXTENSION_KEY.to_string(),
+            extension_name.to_string(),
+        )])))
+    }
+
+    fn return_field(arg_fields: &[FieldRef]) -> FieldRef {
+        let scalars = vec![None; arg_fields.len()];
+        NVL2Func::new()
+            .return_field_from_args(ReturnFieldArgs {
+                arg_fields,
+                scalar_arguments: &scalars,
+            })
+            .unwrap()
+    }
+
+    #[test]
+    fn propagates_metadata_from_value_arguments() {
+        // The first argument is only tested for NULL; its metadata must not
+        // leak into the result even when the value arguments agree.
+        let ret = return_field(&[
+            ext_field("test", DataType::Binary, "unrelated.type"),
+            ext_field("if_non_null", DataType::Binary, "geoarrow.wkb"),
+            ext_field("if_null", DataType::Binary, "geoarrow.wkb"),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Binary);
+        assert_eq!(
+            ret.metadata().get(EXTENSION_KEY).map(String::as_str),
+            Some("geoarrow.wkb")
+        );
+    }
+
+    #[test]
+    fn drops_metadata_when_value_arguments_disagree() {
+        let ret = return_field(&[
+            Arc::new(Field::new("test", DataType::Boolean, true)),
+            ext_field("if_non_null", DataType::Binary, "geoarrow.wkb"),
+            Arc::new(Field::new("if_null", DataType::Binary, true)),
+        ]);
+        assert_eq!(ret.data_type(), &DataType::Binary);
+        assert!(ret.metadata().is_empty());
+    }
+}

datafusion/functions/src/utils.rs

@@ -17,12 +17,13 @@
 
 use arrow::array::{Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray};
 use arrow::compute::try_binary;
-use arrow::datatypes::{DataType, DecimalType};
+use arrow::datatypes::{DataType, DecimalType, Field};
 use arrow::error::ArrowError;
 use datafusion_common::{DataFusionError, Result, ScalarValue};
 use datafusion_expr::ColumnarValue;
 use datafusion_expr::function::Hint;
 use std::cmp::Ordering;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 /// Creates a function to identify the optimal return type of a string function given
@@ -69,6 +70,26 @@ macro_rules! get_optimal_return_type {
     };
 }
 
+/// Returns the field metadata shared by every argument that can contribute a
+/// value to a conditional function's result.
+///
+/// Fields with a `Null` data type (untyped NULL literals) carry no metadata
+/// and are ignored. If the remaining fields disagree on metadata, the result
+/// carries none: propagating one argument's metadata (e.g. an Arrow extension
+/// type name) would claim a type identity for values that other arguments may
+/// supply without it.
+pub(crate) fn unanimous_metadata<'a>(
+    fields: impl Iterator<Item = &'a Field>,
+) -> HashMap<String, String> {
+    let mut candidates = fields.filter(|f| !f.data_type().is_null());
+    match candidates.next() {
+        Some(first) if candidates.all(|f| f.metadata() == first.metadata()) => {
+            first.metadata().clone()
+        }
+        _ => HashMap::new(),
+    }
+}
+
 // `utf8_to_str_type`: returns either a Utf8 or LargeUtf8 based on the input type size.
 get_optimal_return_type!(utf8_to_str_type, DataType::LargeUtf8, DataType::Utf8);