From c3fcddbc2d45be52d173a7faa7a5bff000375e1d Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 2 Oct 2026 03:45:49 -0700 Subject: [PATCH 1/3] fix: align nested array_concat types array_concat planned nested list types from the leaf element type with default inner fields, while execution kept the input fields. LargeList or renamed inner lists failed a type assertion, all-NULL arguments of different depths returned the wrong depth, and non-nullable inner elements panicked. Plan from the dimension-aligned argument types and build the output with the planned type. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/functions-nested/src/concat.rs | 191 +++++++++++------- .../test_files/array/array_concat.slt | 45 ++++- 2 files changed, 156 insertions(+), 80 deletions(-) diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index 4f8c2e7975847..c831d55754df0 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -31,9 +31,7 @@ use arrow::array::{ use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::Result; -use datafusion_common::utils::{ - ListCoercion, base_type, coerced_type_with_base_type_only, offset_span_len, -}; +use datafusion_common::utils::{coerced_fixed_size_list_to_list, offset_span_len}; use datafusion_common::{ cast::as_generic_list_array, exec_err, internal_err, plan_err, @@ -292,45 +290,14 @@ impl ScalarUDFImpl for ArrayConcat { } fn return_type(&self, arg_types: &[DataType]) -> Result { - let mut max_dims = 0; - let mut large_list = false; - let mut element_types = Vec::with_capacity(arg_types.len()); - for arg_type in arg_types { - match arg_type { - DataType::Null | DataType::List(_) | DataType::FixedSizeList(..) => (), - DataType::LargeList(_) => large_list = true, - arg_type => { - return plan_err!("{} does not support type {arg_type}", self.name()); - } - } - - max_dims = max_dims.max(list_ndims(arg_type)); - element_types.push(base_type(arg_type)) - } - - if max_dims == 0 { - Ok(DataType::Null) - } else if let Some(mut return_type) = type_union_resolution(&element_types) { - for _ in 1..max_dims { - return_type = DataType::new_list(return_type, true) - } - - if large_list { - Ok(DataType::new_large_list(return_type, true)) - } else { - Ok(DataType::new_list(return_type, true)) - } - } else { - plan_err!( - "Failed to unify argument types of {}: [{}]", - self.name(), - arg_types.iter().join(", ") - ) - } + array_concat_return_type(self.name(), arg_types) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_concat_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + make_scalar_function(|args: &[ArrayRef]| { + array_concat_with_return_type(args, &return_type) + })(&args.args) } fn aliases(&self) -> &[String] { @@ -339,25 +306,19 @@ impl ScalarUDFImpl for ArrayConcat { fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let return_type = self.return_type(arg_types)?; - let base_type = base_type(&return_type); - let coercion = Some(&ListCoercion::FixedSizedListToList); - // When the return type is a `LargeList`, the outer container of every - // input must be widened to `LargeList` as well. Otherwise - // `array_concat_inner` would later try to downcast a `List` argument - // to `GenericListArray` and fail. - let promote_to_large_list = matches!(return_type, DataType::LargeList(_)); - let arg_types = arg_types.iter().map(|arg_type| { - let coerced = - coerced_type_with_base_type_only(arg_type, &base_type, coercion); - match coerced { - DataType::List(field) if promote_to_large_list => { - DataType::LargeList(field) - } - other => other, - } - }); + let return_ndims = list_ndims(&return_type); - Ok(arg_types.collect()) + Ok(arg_types + .iter() + .map(|arg_type| { + let arg_ndims = list_ndims(arg_type); + if arg_ndims == 0 { + return_type.clone() + } else { + peel_list_type(&return_type, return_ndims - arg_ndims) + } + }) + .collect()) } fn documentation(&self) -> Option<&Documentation> { @@ -365,18 +326,101 @@ impl ScalarUDFImpl for ArrayConcat { } } +fn array_concat_return_type(name: &str, arg_types: &[DataType]) -> Result { + let mut max_dims = 0; + let wrap_with_large_list = arg_types + .iter() + .any(|arg_type| matches!(arg_type, DataType::LargeList(_))); + let mut aligned_types = Vec::with_capacity(arg_types.len()); + + for arg_type in arg_types { + match arg_type { + DataType::Null + | DataType::List(_) + | DataType::FixedSizeList(..) + | DataType::LargeList(_) => (), + arg_type => { + return plan_err!("{name} does not support type {arg_type}"); + } + } + + max_dims = max_dims.max(list_ndims(arg_type)); + } + + if max_dims == 0 { + return Ok(DataType::Null); + } + + for arg_type in arg_types { + aligned_types.push(align_type_dimensions( + arg_type, + max_dims, + wrap_with_large_list, + )); + } + + if let Some(return_type) = type_union_resolution(&aligned_types) { + Ok(return_type) + } else { + plan_err!( + "Failed to unify argument types of {name}: [{}]", + arg_types.iter().join(", ") + ) + } +} + +fn align_type_dimensions( + data_type: &DataType, + max_dims: u64, + wrap_with_large_list: bool, +) -> DataType { + let mut aligned_type = coerced_fixed_size_list_to_list(data_type); + + for _ in list_ndims(data_type)..max_dims { + aligned_type = if wrap_with_large_list { + DataType::new_large_list(aligned_type, true) + } else { + DataType::new_list(aligned_type, true) + }; + } + + aligned_type +} + +fn peel_list_type(data_type: &DataType, levels: u64) -> DataType { + let mut peeled_type = data_type.clone(); + + for _ in 0..levels { + peeled_type = match peeled_type { + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) => field.data_type().clone(), + other => other, + }; + } + + peeled_type +} + pub fn array_concat_inner(args: &[ArrayRef]) -> Result { + let arg_types: Vec = + args.iter().map(|arg| arg.data_type().clone()).collect(); + let return_type = array_concat_return_type("array_concat", &arg_types)?; + array_concat_with_return_type(args, &return_type) +} + +fn array_concat_with_return_type( + args: &[ArrayRef], + return_type: &DataType, +) -> Result { if args.is_empty() { return exec_err!("array_concat expects at least one argument"); } let mut all_null = true; - let mut large_list = false; for arg in args { - match arg.data_type() { - DataType::Null => continue, - DataType::LargeList(_) => large_list = true, - _ => (), + if arg.data_type() == &DataType::Null { + continue; } if arg.null_count() < arg.len() { all_null = false; @@ -384,21 +428,17 @@ pub fn array_concat_inner(args: &[ArrayRef]) -> Result { } if all_null { - // Return a null array with the same type as the first non-null-type argument - let return_type = args - .iter() - .map(|arg| arg.data_type()) - .find_or_first(|d| !d.is_null()) - .unwrap(); // Safe because args is non-empty - Ok(arrow::array::make_array(ArrayData::new_null( return_type, args[0].len(), ))) - } else if large_list { - concat_internal::(args, None) } else { - concat_internal::(args, None) + let field = list_inner_field("array_concat", return_type)?; + match return_type { + DataType::List(_) => concat_internal::(args, Some(&field)), + DataType::LargeList(_) => concat_internal::(args, Some(&field)), + arg_type => exec_err!("array_concat does not support type {arg_type}"), + } } } @@ -419,11 +459,8 @@ fn append_prepend_return_type( /// Concatenates the list arrays in `args` row-wise. /// -/// `field` is the list field the output must carry. `array_concat` passes `None` -/// because its `return_type` derives a fresh field from the unified element -/// types, which is what deriving the field from the aligned inputs reproduces. -/// `array_append` / `array_prepend` promise their input's field verbatim and so -/// must pass it in explicitly. +/// `field` is the list field the output must carry when it differs from the +/// aligned inputs' first field. fn concat_internal( args: &[ArrayRef], field: Option<&FieldRef>, diff --git a/datafusion/sqllogictest/test_files/array/array_concat.slt b/datafusion/sqllogictest/test_files/array/array_concat.slt index 5b7985ef6d194..f548861f6e33b 100644 --- a/datafusion/sqllogictest/test_files/array/array_concat.slt +++ b/datafusion/sqllogictest/test_files/array/array_concat.slt @@ -419,9 +419,8 @@ select array_concat(make_array(column3), column1, column2) from arrays_values_v2 [NULL, 11, 12] [NULL] -# array_concat derives a fresh return type from the unified element types rather -# than cloning an input's, so its output field is always the default nullable -# `item` regardless of the inputs' inner fields +# For one-dimensional inputs, array_concat uses the unified element type and a +# default nullable item field rather than cloning the input field. query ?T select array_concat(arrow_cast(column1, 'List(non-null Int64)'), make_array(3)), @@ -438,5 +437,45 @@ from values (make_array(1, 2)); ---- [1, 2, 3] List(Int64) +# array_concat preserves nested list fields in the planned and executed type +query ?T +select + array_concat(a, a), + arrow_typeof(array_concat(a, a)) +from ( + select make_array(arrow_cast([1, 2], 'LargeList(Int64)')) as a +); +---- +[[1, 2], [1, 2]] List(LargeList(Int64)) + +query ?T +select + array_concat(b, b), + arrow_typeof(array_concat(b, b)) +from ( + select arrow_cast([[1]], 'List(List(Int64, field: ''element''), field: ''element'')') as b +); +---- +[[1], [1]] List(List(Int64, field: 'element'), field: 'element') + +# array_concat aligns all-null arguments to the deepest list type +query ?T +select + array_concat(arrow_cast(NULL, 'List(Int64)'), arrow_cast(NULL, 'List(List(Int64))')), + arrow_typeof(array_concat(arrow_cast(NULL, 'List(Int64)'), arrow_cast(NULL, 'List(List(Int64))'))); +---- +NULL List(List(Int64)) + +# array_concat widens nested element nullability before concatenating +query ?T +select + array_concat(c, make_array(2)), + arrow_typeof(array_concat(c, make_array(2))) +from ( + select arrow_cast(make_array(make_array(1)), 'List(List(non-null Int64))') as c +); +---- +[[1], [2]] List(List(Int64)) + include ./cleanup.slt.part From 07cf61add4cc64feaf1a894597fc13f9c6956648 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Fri, 2 Oct 2026 06:16:59 -0700 Subject: [PATCH 2/3] fix: keep the default item field for one-dimensional array_concat One-dimensional inputs that share a named or non-nullable field kept that field in the result, which changed the result type from main. Plan one-dimensional results from the unified element type with the default nullable field again, and pin it in array_concat.slt. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/functions-nested/src/concat.rs | 19 ++++++++++++++++++- .../test_files/array/array_concat.slt | 12 ++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index c831d55754df0..8263f9e78141a 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -31,7 +31,9 @@ use arrow::array::{ use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::Result; -use datafusion_common::utils::{coerced_fixed_size_list_to_list, offset_span_len}; +use datafusion_common::utils::{ + base_type, coerced_fixed_size_list_to_list, offset_span_len, +}; use datafusion_common::{ cast::as_generic_list_array, exec_err, internal_err, plan_err, @@ -351,6 +353,21 @@ fn array_concat_return_type(name: &str, arg_types: &[DataType]) -> Result = arg_types.iter().map(base_type).collect(); + let Some(element_type) = type_union_resolution(&element_types) else { + return plan_err!( + "Failed to unify argument types of {name}: [{}]", + arg_types.iter().join(", ") + ); + }; + return Ok(if wrap_with_large_list { + DataType::new_large_list(element_type, true) + } else { + DataType::new_list(element_type, true) + }); + } + for arg_type in arg_types { aligned_types.push(align_type_dimensions( arg_type, diff --git a/datafusion/sqllogictest/test_files/array/array_concat.slt b/datafusion/sqllogictest/test_files/array/array_concat.slt index f548861f6e33b..bcf1cad846c2e 100644 --- a/datafusion/sqllogictest/test_files/array/array_concat.slt +++ b/datafusion/sqllogictest/test_files/array/array_concat.slt @@ -437,6 +437,18 @@ from values (make_array(1, 2)); ---- [1, 2, 3] List(Int64) +# One-dimensional inputs that share a named or non-nullable field still get the +# default nullable item field. +query ?T +select + array_concat(a, a), + arrow_typeof(array_concat(a, a)) +from ( + select arrow_cast([1], 'List(non-null Int64, field: ''element'')') as a +); +---- +[1, 1] List(Int64) + # array_concat preserves nested list fields in the planned and executed type query ?T select From 5b91b532420990dc31d790fe2e301959be2b1f80 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 22 Oct 2026 22:00:15 -0700 Subject: [PATCH 3/3] docs: clarify concat field sources Signed-off-by: 1fanwang <1fannnw@gmail.com> --- datafusion/functions-nested/src/concat.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index 8263f9e78141a..85204042b9349 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -476,8 +476,9 @@ fn append_prepend_return_type( /// Concatenates the list arrays in `args` row-wise. /// -/// `field` is the list field the output must carry when it differs from the -/// aligned inputs' first field. +/// `field` is the list field the output must carry. `array_concat` passes the +/// inner field of its planned return type, and `array_append` / `array_prepend` +/// pass their input's field. `None` derives the field from the aligned inputs. fn concat_internal( args: &[ArrayRef], field: Option<&FieldRef>,