diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index 9cf3536443e6..b48ab52f2db1 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/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") + ); + } +} diff --git a/datafusion/functions/src/core/nvl2.rs b/datafusion/functions/src/core/nvl2.rs index d68296d9b862..612b2b32438e 100644 --- a/datafusion/functions/src/core/nvl2.rs +++ b/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 { @@ -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()); + } +} diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index 52ce0bd72cac..2a893cb6c228 100644 --- a/datafusion/functions/src/utils.rs +++ b/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, +) -> HashMap { + 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);