diff --git a/datafusion/common/src/metadata.rs b/datafusion/common/src/metadata.rs index 34242e9b52757..2b93141a4c2e9 100644 --- a/datafusion/common/src/metadata.rs +++ b/datafusion/common/src/metadata.rs @@ -18,6 +18,7 @@ use std::{collections::BTreeMap, sync::Arc}; use arrow::datatypes::{DataType, Field, FieldRef, Metadata}; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use hashbrown::HashMap; use crate::{DataFusionError, ScalarValue, error::_plan_err}; @@ -117,6 +118,75 @@ pub fn check_metadata_with_storage_equal( Ok(()) } +/// Compute the field metadata to propagate for an expression whose result +/// may come from any one of several inputs coerced to a common type, such as +/// `COALESCE`, `NVL2`, or `CASE WHEN`. +/// +/// Metadata is merged across the fields whose data type is not `Null` (an +/// untyped NULL literal carries no metadata and cannot contribute a typed +/// value), with earlier fields taking precedence for repeated keys. Unrelated +/// keys are not required to agree: metadata that arrived with the data (e.g. +/// from an Arrow file or an embedded Arrow schema in a Parquet file) must not +/// cause an extension type to be silently dropped from the result. +/// +/// The Arrow extension type keys are handled as a unit, like in +/// [`Expr::Cast`] output field derivation: +/// - Differing values of `ARROW:extension:name` produce a planning error; +/// the inputs are of different logical types even though they were coerced +/// to a common storage type. +/// - `ARROW:extension:metadata` is taken from the field that supplied the +/// extension name, never combined from another field. Byte-wise differences +/// in the extension metadata of fields that agree on the extension name are +/// tolerated, since extension parameters are often JSON-encoded, where byte +/// equality is too strict. +/// +/// [`Expr::Cast`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html#variant.Cast +pub fn coerced_fields_metadata<'a>( + fields: impl IntoIterator, +) -> Result { + let mut merged: BTreeMap = BTreeMap::new(); + let mut extension: Option<(String, Option)> = None; + + for field in fields { + if field.data_type().is_null() { + continue; + } + let metadata = FieldMetadata::new_from_field(field); + if let Some(name) = metadata.inner().get(EXTENSION_TYPE_NAME_KEY) { + match &extension { + Some((first_name, _)) if first_name != name => { + return _plan_err!( + "Cannot merge field metadata with differing values of \ + '{EXTENSION_TYPE_NAME_KEY}': '{first_name}' and '{name}'" + ); + } + Some(_) => {} + None => { + extension = Some(( + name.clone(), + metadata.inner().get(EXTENSION_TYPE_METADATA_KEY).cloned(), + )); + } + } + } + for (key, value) in metadata.inner() { + if key == EXTENSION_TYPE_NAME_KEY || key == EXTENSION_TYPE_METADATA_KEY { + continue; + } + merged.entry(key.clone()).or_insert_with(|| value.clone()); + } + } + + if let Some((name, extension_metadata)) = extension { + merged.insert(EXTENSION_TYPE_NAME_KEY.to_string(), name); + if let Some(extension_metadata) = extension_metadata { + merged.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), extension_metadata); + } + } + + Ok(FieldMetadata::from(merged)) +} + /// Given a data type represented by storage and optional metadata, generate /// a user-facing string /// @@ -392,3 +462,98 @@ impl From<&HashMap> for FieldMetadata { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn coerced_fields_metadata_merges_with_earlier_precedence() { + let a = Field::new("a", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([ + ("k".to_string(), "v1".to_string()), + ("only_a".to_string(), "a".to_string()), + ]), + ); + let b = Field::new("b", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([ + ("k".to_string(), "v2".to_string()), + ("only_b".to_string(), "b".to_string()), + ]), + ); + let result = coerced_fields_metadata([&a, &b]).unwrap(); + assert_eq!(result.inner().get("k").map(String::as_str), Some("v1")); + assert_eq!(result.inner().get("only_a").map(String::as_str), Some("a")); + assert_eq!(result.inner().get("only_b").map(String::as_str), Some("b")); + } + + #[test] + fn coerced_fields_metadata_skips_null_typed_fields() { + let null = Field::new("null", DataType::Null, true); + let b = Field::new("b", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([("k".to_string(), "v".to_string())]), + ); + let result = coerced_fields_metadata([&null, &b]).unwrap(); + assert_eq!(result.inner().get("k").map(String::as_str), Some("v")); + } + + #[test] + fn coerced_fields_metadata_empty_when_no_field_has_any() { + let a = Field::new("a", DataType::Binary, true); + let null = Field::new("null", DataType::Null, true); + assert!(coerced_fields_metadata([&a, &null]).unwrap().is_empty()); + } + + #[test] + fn coerced_fields_metadata_errors_on_differing_extension_names() { + let a = Field::new("a", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + )]), + ); + let b = Field::new("b", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]), + ); + let err = coerced_fields_metadata([&a, &b]).unwrap_err(); + assert!(err.to_string().contains("arrow.json"), "{err}"); + assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); + } + + #[test] + fn coerced_fields_metadata_keeps_extension_keys_paired() { + // The extension metadata travels with the field that supplied the + // extension name: byte-wise differences for the same name are + // tolerated, and a bare extension name never picks up another + // field's parameters. + let a = Field::new("a", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "ext.type".to_string(), + )]), + ); + let b = Field::new("b", DataType::Binary, true).with_metadata( + std::collections::HashMap::from([ + (EXTENSION_TYPE_NAME_KEY.to_string(), "ext.type".to_string()), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "{\"param\": 1}".to_string(), + ), + ]), + ); + let result = coerced_fields_metadata([&a, &b]).unwrap(); + assert_eq!( + result + .inner() + .get(EXTENSION_TYPE_NAME_KEY) + .map(String::as_str), + Some("ext.type") + ); + // "a" supplied the extension name first and carried no parameters; + // "b"'s parameters are not grafted onto it. + assert_eq!(result.inner().get(EXTENSION_TYPE_METADATA_KEY), None); + } +} diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index 9cf3536443e6c..10733d0ff0e86 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/datafusion/functions/src/core/coalesce.rs @@ -16,6 +16,7 @@ // under the License. use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::metadata::coerced_fields_metadata; use datafusion_common::{Result, exec_err, internal_err, plan_err}; use datafusion_expr::binary::try_type_union_resolution; use datafusion_expr::conditional_expressions::CaseBuilder; @@ -86,7 +87,14 @@ 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) from the + // arguments that can supply a value; see `coerced_fields_metadata` + // for the merge semantics. + let metadata = + coerced_fields_metadata(args.arg_fields.iter().map(|f| f.as_ref()))?; + Ok(metadata + .add_to_field(Field::new(self.name(), return_type, nullable)) + .into()) } fn simplify( @@ -146,3 +154,142 @@ 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"; + + #[test] + fn propagates_metadata_when_arguments_agree() { + let fields = [ + Arc::new(Field::new("a", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + Arc::new(Field::new("b", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + ]; + let ret = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap(); + 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 is skipped when picking the metadata. + let fields = [ + Arc::new(Field::new("null", DataType::Null, true)), + Arc::new(Field::new("b", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + ]; + let ret = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap(); + assert_eq!(ret.data_type(), &DataType::Binary); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); + } + + #[test] + fn errors_on_conflicting_extension_types() { + let fields = [ + Arc::new(Field::new("a", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + Arc::new(Field::new("b", DataType::Binary, true).with_metadata( + HashMap::from([( + EXTENSION_KEY.to_string(), + "unrelated.type".to_string(), + )]), + )), + ]; + let err = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap_err(); + assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); + assert!(err.to_string().contains("unrelated.type"), "{err}"); + } + + #[test] + fn plain_argument_does_not_block_later_metadata() { + let fields = [ + Arc::new(Field::new("plain", DataType::Binary, true)), + Arc::new(Field::new("b", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + ]; + let ret = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap(); + assert_eq!(ret.data_type(), &DataType::Binary); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); + } + + #[test] + fn no_metadata_when_all_arguments_are_null() { + let fields = [ + Arc::new(Field::new("a", DataType::Null, true)), + Arc::new(Field::new("b", DataType::Null, true)), + ]; + let ret = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap(); + assert_eq!(ret.data_type(), &DataType::Null); + assert!(ret.metadata().is_empty()); + } + + #[test] + fn nullability_is_unchanged_by_metadata_propagation() { + let fields = [ + Arc::new(Field::new("a", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + Arc::new(Field::new("b", DataType::Binary, false).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + ]; + let ret = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None], + }) + .unwrap(); + 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 671e1d74f5f89..e0bd860f68f21 100644 --- a/datafusion/functions/src/core/nvl2.rs +++ b/datafusion/functions/src/core/nvl2.rs @@ -16,6 +16,7 @@ // under the License. use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::metadata::coerced_fields_metadata; use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args}; use datafusion_expr::{ ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, @@ -94,7 +95,14 @@ 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) from them. The first + // argument is only tested for NULL and does not contribute a value. + let metadata = + coerced_fields_metadata(args.arg_fields[1..3].iter().map(|f| f.as_ref()))?; + Ok(metadata + .add_to_field(Field::new(self.name(), return_type, nullable)) + .into()) } fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { @@ -140,3 +148,77 @@ 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"; + + #[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 fields = [ + Arc::new(Field::new("test", DataType::Binary, true).with_metadata( + HashMap::from([( + EXTENSION_KEY.to_string(), + "unrelated.type".to_string(), + )]), + )), + Arc::new( + Field::new("if_non_null", DataType::Binary, true).with_metadata( + HashMap::from([( + EXTENSION_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]), + ), + ), + Arc::new(Field::new("if_null", DataType::Binary, true).with_metadata( + HashMap::from([(EXTENSION_KEY.to_string(), "geoarrow.wkb".to_string())]), + )), + ]; + let ret = NVL2Func::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None, None], + }) + .unwrap(); + assert_eq!(ret.data_type(), &DataType::Binary); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); + } + + #[test] + fn errors_when_value_arguments_have_conflicting_extension_types() { + let fields = [ + Arc::new(Field::new("test", DataType::Boolean, true)), + Arc::new( + Field::new("if_non_null", DataType::Binary, true).with_metadata( + HashMap::from([( + EXTENSION_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]), + ), + ), + Arc::new(Field::new("if_null", DataType::Binary, true).with_metadata( + HashMap::from([( + EXTENSION_KEY.to_string(), + "unrelated.type".to_string(), + )]), + )), + ]; + let err = NVL2Func::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &fields, + scalar_arguments: &[None, None, None], + }) + .unwrap_err(); + assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); + assert!(err.to_string().contains("unrelated.type"), "{err}"); + } +}