From 1ea88c0cbd8b5b919b3b958974d86d77beb6f78a Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 1 Sep 2026 13:54:20 -0700 Subject: [PATCH 1/4] fix: propagate field metadata through coalesce, nvl, and nvl2 These functions built their return Field from argument data types only, dropping field metadata (e.g. Arrow extension types) from the planned schema. Attach the metadata shared by every value-contributing argument instead. Part of #24860. --- datafusion/functions/src/core/coalesce.rs | 102 +++++++++++++++++++++- datafusion/functions/src/core/nvl2.rs | 64 +++++++++++++- datafusion/functions/src/utils.rs | 23 ++++- 3 files changed, 186 insertions(+), 3 deletions(-) diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index 9cf3536443e6c..b48ab52f2db11 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 d68296d9b862b..612b2b32438e6 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 52ce0bd72cac5..2a893cb6c228d 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); From 33fe79d4639b6e78823f830bd9a3a028413324f1 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Wed, 2 Sep 2026 15:07:48 -0700 Subject: [PATCH 2/4] Address review: propagate metadata leniently, move helper to datafusion-common Requiring the arguments to agree on metadata was too strict: extension types with JSON-encoded parameters can differ byte-wise, and unrelated metadata that arrived with the data should not drop an extension type from the result. Propagate the first non-empty metadata among the value-contributing arguments instead, and move the helper next to the related utilities in datafusion_common::metadata so the upcoming CASE fix in datafusion-expr can reuse it. --- datafusion/common/src/metadata.rs | 69 +++++++++++++++++++++++ datafusion/functions/src/core/coalesce.rs | 36 +++++++++--- datafusion/functions/src/core/nvl2.rs | 22 ++++---- datafusion/functions/src/utils.rs | 23 +------- 4 files changed, 109 insertions(+), 41 deletions(-) diff --git a/datafusion/common/src/metadata.rs b/datafusion/common/src/metadata.rs index d6d8fb7b0ed0c..5a477c1885c53 100644 --- a/datafusion/common/src/metadata.rs +++ b/datafusion/common/src/metadata.rs @@ -123,6 +123,36 @@ 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`. +/// +/// Returns the first non-empty metadata among the fields whose data type is +/// not `Null` (an untyped NULL literal carries no metadata and cannot +/// contribute a typed value), or empty metadata if there is none. +/// +/// The inputs are deliberately not required to agree on their metadata: +/// byte-wise comparison is too strict for extension types whose parameters +/// are JSON-encoded, and unrelated 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. Engines +/// that want stricter behavior (e.g. rejecting conflicting extension types) +/// can enforce it with an analyzer or optimizer rule, which can only observe +/// the conflict if planning preserves the metadata in the first place. +pub fn coerced_fields_metadata<'a>( + mut fields: impl Iterator, +) -> FieldMetadata { + fields + .find_map(|f| { + if f.data_type().is_null() { + return None; + } + let metadata = FieldMetadata::new_from_field(f); + (!metadata.is_empty()).then_some(metadata) + }) + .unwrap_or_default() +} + /// Given a data type represented by storage and optional metadata, generate /// a user-facing string /// @@ -382,3 +412,42 @@ impl From<&HashMap> for FieldMetadata { Self::new(inner) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn field_with(name: &str, dt: DataType, pairs: &[(&str, &str)]) -> Field { + Field::new(name, dt, true).with_metadata( + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ) + } + + #[test] + fn coerced_fields_metadata_returns_first_non_empty() { + let a = field_with("a", DataType::Binary, &[]); + let b = field_with("b", DataType::Binary, &[("k", "v1")]); + let c = field_with("c", DataType::Binary, &[("k", "v2")]); + let result = coerced_fields_metadata([&a, &b, &c].into_iter()); + assert_eq!(result.inner().get("k").map(String::as_str), Some("v1")); + } + + #[test] + fn coerced_fields_metadata_skips_null_typed_fields() { + let null = field_with("null", DataType::Null, &[]); + let b = field_with("b", DataType::Binary, &[("k", "v")]); + let result = coerced_fields_metadata([&null, &b].into_iter()); + 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_with("a", DataType::Binary, &[]); + let null = field_with("null", DataType::Null, &[]); + let result = coerced_fields_metadata([&a, &null].into_iter()); + assert!(result.is_empty()); + } +} diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index b48ab52f2db11..20276fee27d76 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/datafusion/functions/src/core/coalesce.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::unanimous_metadata; 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; @@ -87,11 +87,13 @@ impl ScalarUDFImpl for CoalesceFunc { .find_or_first(|d| !d.is_null()) .unwrap() .clone(); - // 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) + // Propagate field metadata (e.g. Arrow extension types) from the + // arguments that can supply a value; see `coerced_fields_metadata` + // for why the arguments are not required to agree on it. + 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()) } @@ -199,7 +201,7 @@ mod tests { #[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. + // a typed value either, so it is skipped when picking the metadata. let ret = return_field(&[ field("null", DataType::Null, true), ext_field("b", DataType::Binary, "geoarrow.wkb"), @@ -212,13 +214,29 @@ mod tests { } #[test] - fn drops_metadata_when_arguments_disagree() { + fn first_metadata_wins_when_arguments_disagree() { let ret = return_field(&[ ext_field("a", DataType::Binary, "geoarrow.wkb"), + ext_field("b", DataType::Binary, "unrelated.type"), + ]); + assert_eq!(ret.data_type(), &DataType::Binary); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); + } + + #[test] + fn plain_argument_does_not_block_later_metadata() { + let ret = return_field(&[ field("plain", DataType::Binary, true), + ext_field("b", DataType::Binary, "geoarrow.wkb"), ]); assert_eq!(ret.data_type(), &DataType::Binary); - assert!(ret.metadata().is_empty()); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); } #[test] diff --git a/datafusion/functions/src/core/nvl2.rs b/datafusion/functions/src/core/nvl2.rs index 612b2b32438e6..cc131757110cb 100644 --- a/datafusion/functions/src/core/nvl2.rs +++ b/datafusion/functions/src/core/nvl2.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::unanimous_metadata; use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::metadata::coerced_fields_metadata; use datafusion_common::{Result, internal_err, utils::take_function_args}; use datafusion_expr::{ ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, @@ -96,13 +96,12 @@ impl ScalarUDFImpl for NVL2Func { args.arg_fields[1].is_nullable() || args.arg_fields[2].is_nullable(); let return_type = args.arg_fields[1].data_type().clone(); // 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. + // 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 = - unanimous_metadata(args.arg_fields[1..3].iter().map(|f| f.as_ref())); - Ok(Field::new(self.name(), return_type, nullable) - .with_metadata(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()) } @@ -197,13 +196,16 @@ mod tests { } #[test] - fn drops_metadata_when_value_arguments_disagree() { + fn first_value_argument_metadata_wins_when_they_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)), + ext_field("if_null", DataType::Binary, "unrelated.type"), ]); assert_eq!(ret.data_type(), &DataType::Binary); - assert!(ret.metadata().is_empty()); + assert_eq!( + ret.metadata().get(EXTENSION_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); } } diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index 2a893cb6c228d..52ce0bd72cac5 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -17,13 +17,12 @@ use arrow::array::{Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray}; use arrow::compute::try_binary; -use arrow::datatypes::{DataType, DecimalType, Field}; +use arrow::datatypes::{DataType, DecimalType}; 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 @@ -70,26 +69,6 @@ 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); From e76a9b82b4e8ff80fad9572096e627cee6da18b0 Mon Sep 17 00:00:00 2001 From: jameswillis Date: Mon, 28 Sep 2026 15:30:40 -0700 Subject: [PATCH 3/4] Address review: merge metadata, error on conflicting extension names Adopt the suggested compromise: merge metadata across the value-contributing arguments (earlier arguments win for repeated keys) instead of taking only the first non-empty metadata, but refuse differing values of 'ARROW:extension:name' with a planning error, keeping the extension name and its parameters paired. Also accept impl IntoIterator and inline the field construction in the datafusion-common tests. --- datafusion/common/src/metadata.rs | 179 +++++++++++++++++----- datafusion/functions/src/core/coalesce.rs | 24 +-- datafusion/functions/src/core/nvl2.rs | 22 +-- 3 files changed, 165 insertions(+), 60 deletions(-) diff --git a/datafusion/common/src/metadata.rs b/datafusion/common/src/metadata.rs index ad65fc197bf7d..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}; @@ -121,30 +122,69 @@ pub fn check_metadata_with_storage_equal( /// may come from any one of several inputs coerced to a common type, such as /// `COALESCE`, `NVL2`, or `CASE WHEN`. /// -/// Returns the first non-empty metadata among the fields whose data type is -/// not `Null` (an untyped NULL literal carries no metadata and cannot -/// contribute a typed value), or empty metadata if there is none. -/// -/// The inputs are deliberately not required to agree on their metadata: -/// byte-wise comparison is too strict for extension types whose parameters -/// are JSON-encoded, and unrelated metadata that arrived with the data (e.g. +/// 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. Engines -/// that want stricter behavior (e.g. rejecting conflicting extension types) -/// can enforce it with an analyzer or optimizer rule, which can only observe -/// the conflict if planning preserves the metadata in the first place. +/// 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>( - mut fields: impl Iterator, -) -> FieldMetadata { - fields - .find_map(|f| { - if f.data_type().is_null() { - return None; + 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(), + )); + } } - let metadata = FieldMetadata::new_from_field(f); - (!metadata.is_empty()).then_some(metadata) - }) - .unwrap_or_default() + } + 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 @@ -427,36 +467,93 @@ impl From<&HashMap> for FieldMetadata { mod tests { use super::*; - fn field_with(name: &str, dt: DataType, pairs: &[(&str, &str)]) -> Field { - let metadata: std::collections::HashMap = pairs - .iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect(); - Field::new(name, dt, true).with_metadata(metadata) - } - #[test] - fn coerced_fields_metadata_returns_first_non_empty() { - let a = field_with("a", DataType::Binary, &[]); - let b = field_with("b", DataType::Binary, &[("k", "v1")]); - let c = field_with("c", DataType::Binary, &[("k", "v2")]); - let result = coerced_fields_metadata([&a, &b, &c].into_iter()); + 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_with("null", DataType::Null, &[]); - let b = field_with("b", DataType::Binary, &[("k", "v")]); - let result = coerced_fields_metadata([&null, &b].into_iter()); + 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_with("a", DataType::Binary, &[]); - let null = field_with("null", DataType::Null, &[]); - let result = coerced_fields_metadata([&a, &null].into_iter()); - assert!(result.is_empty()); + 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 20276fee27d76..6e98301a2f58a 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/datafusion/functions/src/core/coalesce.rs @@ -89,9 +89,9 @@ impl ScalarUDFImpl for CoalesceFunc { .clone(); // Propagate field metadata (e.g. Arrow extension types) from the // arguments that can supply a value; see `coerced_fields_metadata` - // for why the arguments are not required to agree on it. + // for the merge semantics. let metadata = - coerced_fields_metadata(args.arg_fields.iter().map(|f| f.as_ref())); + 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()) @@ -214,16 +214,20 @@ mod tests { } #[test] - fn first_metadata_wins_when_arguments_disagree() { - let ret = return_field(&[ + fn errors_on_conflicting_extension_types() { + let arg_fields = [ ext_field("a", DataType::Binary, "geoarrow.wkb"), ext_field("b", DataType::Binary, "unrelated.type"), - ]); - assert_eq!(ret.data_type(), &DataType::Binary); - assert_eq!( - ret.metadata().get(EXTENSION_KEY).map(String::as_str), - Some("geoarrow.wkb") - ); + ]; + let scalars = vec![None; arg_fields.len()]; + let err = CoalesceFunc::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalars, + }) + .unwrap_err(); + assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); + assert!(err.to_string().contains("unrelated.type"), "{err}"); } #[test] diff --git a/datafusion/functions/src/core/nvl2.rs b/datafusion/functions/src/core/nvl2.rs index cc131757110cb..9a554b7c868c8 100644 --- a/datafusion/functions/src/core/nvl2.rs +++ b/datafusion/functions/src/core/nvl2.rs @@ -99,7 +99,7 @@ impl ScalarUDFImpl for NVL2Func { // 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())); + 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()) @@ -196,16 +196,20 @@ mod tests { } #[test] - fn first_value_argument_metadata_wins_when_they_disagree() { - let ret = return_field(&[ + fn errors_when_value_arguments_have_conflicting_extension_types() { + let arg_fields = [ Arc::new(Field::new("test", DataType::Boolean, true)), ext_field("if_non_null", DataType::Binary, "geoarrow.wkb"), ext_field("if_null", DataType::Binary, "unrelated.type"), - ]); - assert_eq!(ret.data_type(), &DataType::Binary); - assert_eq!( - ret.metadata().get(EXTENSION_KEY).map(String::as_str), - Some("geoarrow.wkb") - ); + ]; + let scalars = vec![None; arg_fields.len()]; + let err = NVL2Func::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalars, + }) + .unwrap_err(); + assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); + assert!(err.to_string().contains("unrelated.type"), "{err}"); } } From f70a43ce10e69ab4cb986a728dfa95b6d09bf93e Mon Sep 17 00:00:00 2001 From: jameswillis Date: Tue, 27 Oct 2026 11:14:38 -0700 Subject: [PATCH 4/4] Address review: inline test field construction in coalesce and nvl2 Drop the field/ext_field/return_field test helpers in favor of Field::new().with_metadata() and direct return_field_from_args calls, matching the datafusion-common tests in this PR and the existing tests in with_metadata.rs. --- datafusion/functions/src/core/coalesce.rs | 117 +++++++++++++--------- datafusion/functions/src/core/nvl2.rs | 70 +++++++------ 2 files changed, 113 insertions(+), 74 deletions(-) diff --git a/datafusion/functions/src/core/coalesce.rs b/datafusion/functions/src/core/coalesce.rs index 6e98301a2f58a..10733d0ff0e86 100644 --- a/datafusion/functions/src/core/coalesce.rs +++ b/datafusion/functions/src/core/coalesce.rs @@ -163,33 +163,22 @@ mod tests { 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"), - ]); + 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!( @@ -202,10 +191,18 @@ mod tests { 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 ret = return_field(&[ - field("null", DataType::Null, true), - ext_field("b", DataType::Binary, "geoarrow.wkb"), - ]); + 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), @@ -215,15 +212,21 @@ mod tests { #[test] fn errors_on_conflicting_extension_types() { - let arg_fields = [ - ext_field("a", DataType::Binary, "geoarrow.wkb"), - ext_field("b", DataType::Binary, "unrelated.type"), + 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 scalars = vec![None; arg_fields.len()]; let err = CoalesceFunc::new() .return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &scalars, + arg_fields: &fields, + scalar_arguments: &[None, None], }) .unwrap_err(); assert!(err.to_string().contains("geoarrow.wkb"), "{err}"); @@ -232,10 +235,18 @@ mod tests { #[test] fn plain_argument_does_not_block_later_metadata() { - let ret = return_field(&[ - field("plain", DataType::Binary, true), - ext_field("b", DataType::Binary, "geoarrow.wkb"), - ]); + 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), @@ -245,22 +256,36 @@ mod tests { #[test] fn no_metadata_when_all_arguments_are_null() { - let ret = return_field(&[ - field("a", DataType::Null, true), - field("b", DataType::Null, true), - ]); + 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 ret = return_field(&[ - ext_field("a", DataType::Binary, "geoarrow.wkb"), + 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), diff --git a/datafusion/functions/src/core/nvl2.rs b/datafusion/functions/src/core/nvl2.rs index a08d90618b2ea..e0bd860f68f21 100644 --- a/datafusion/functions/src/core/nvl2.rs +++ b/datafusion/functions/src/core/nvl2.rs @@ -157,32 +157,35 @@ mod tests { 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"), - ]); + 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), @@ -192,16 +195,27 @@ mod tests { #[test] fn errors_when_value_arguments_have_conflicting_extension_types() { - let arg_fields = [ + let fields = [ Arc::new(Field::new("test", DataType::Boolean, true)), - ext_field("if_non_null", DataType::Binary, "geoarrow.wkb"), - ext_field("if_null", DataType::Binary, "unrelated.type"), + 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 scalars = vec![None; arg_fields.len()]; let err = NVL2Func::new() .return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &scalars, + arg_fields: &fields, + scalar_arguments: &[None, None, None], }) .unwrap_err(); assert!(err.to_string().contains("geoarrow.wkb"), "{err}");