From a5874174d3c77a9f74d6ae4ddb7df4bcf19cd267 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 14 Aug 2026 17:07:04 +0000 Subject: [PATCH] fix: preserve the input list's inner field in array_append/prepend/replace (#24365) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #24347. ## Rationale for this change `array_append`, `array_prepend`, `array_replace`, `array_replace_n` and `array_replace_all` promise the input list type verbatim — inner field name, nullability and metadata included — but their kernels rebuilt the output's inner field from scratch with `Field::new_list_field(..., true)`. On debug builds this trips the return-type assertion from #17515; on release builds it silently yields a batch whose inner field disagrees with the schema the planner recorded. This is the other half of #24341, whose `array_slice` part was fixed in #24345. The field-name symptom is a 55.0.0 regression from the same commit (5b228570, #20945); the non-nullable symptom is not a regression. Unlike `array_slice`, threading the input's field through is not sufficient here: the appended, prepended or replacement element can itself be null, so a promise cloned from a `List(non-null T)` input is wrong at the source and arrow rejects the array with `Non-nullable field of ListArray cannot contain nulls`. ## What changes are included in this PR? Each of the five functions now implements `return_field_from_args`, carrying the input field's name and metadata through while widening `nullable` when the new element's argument is nullable. The kernels build their output from `args.return_field` instead of deriving a field of their own, so promise and payload come from a single source. Nullability is therefore only widened when the result can genuinely contain a null: ```sql array_append(List(non-null Int64), 3) -> List(non-null Int64) array_append(List(non-null Int64), NULL) -> List(Int64) ``` Two paths beyond those listed in the issue turned out to have the same defect and are fixed too: `array_append` / `array_prepend` with a **nested** value type (which delegates to `concat_internal`), and the `LargeList` variants of all five. Since these five now implement `return_field_from_args`, their `return_type` becomes unreachable and returns `internal_err!("return_field_from_args should be used instead")`, matching the guidance on `ScalarUDFImpl::return_type` and the existing convention in `remove.rs` and `map_values.rs`. Two small cleanups while in here: - The `List`/`LargeList` inner-field extraction added to `general_array_slice` by #24345 is now shared as `utils::list_inner_field`, used by all three files. Its error text is unchanged. The `ListView` variant in `general_list_view_array_slice` is deliberately left alone — folding all four variants into one helper would let `general_array_slice` silently accept a `ListView` that its match currently rejects. - `array_concat` is **not** affected and its behaviour is unchanged: it derives a fresh return type via `type_union_resolution` rather than cloning an input's, so it keeps passing `None` to `concat_internal` and deriving the field from the aligned inputs. Behaviour for `Null`-typed array arguments is unchanged in all five functions (`array_replace*` return `Null`, `array_append` / `array_prepend` return `List(element)`). ## Are these changes tested? Yes — 14 new SLT tests across `array_append.slt`, `array_prepend.slt`, `array_replace.slt`, plus two guard cases in `array_concat.slt` pinning down that it is unaffected. `arrow_cast` can express a named inner field (`'List(Int64, field: ''element'')'`), so these reproduce the field-name half of the bug without needing the Spark dialect. Coverage: inner field name preserved, non-nullable inner field preserved, nullability widened only when the new element is nullable (both literal `NULL` and a nullable column), `LargeList`, nested value types, the `max <= 0` short circuit, and a `NULL` `max`. Every one of these queries fails on `main` with the return-type assertion. Also run: `cargo clippy --all-targets --all-features -- -D warnings`, the full sqllogictest suite, and the extended workspace test suite (68 test binaries, 0 failures). The `array_replace` and `array_concat` benchmarks show no regression against `main`. ## Are there any user-facing changes? The five functions now return the inner field they promise instead of a rebuilt one, which is the bug fix. As a consequence, appending or replacing with a nullable element widens the declared inner nullability of the result (`List(non-null Int64)` -> `List(Int64)`), which is required for the result to be representable at all. No breaking changes to public APIs. --- datafusion/functions-nested/src/concat.rs | 140 +++++++++++----- datafusion/functions-nested/src/extract.rs | 11 +- datafusion/functions-nested/src/replace.rs | 157 +++++++++++++----- datafusion/functions-nested/src/utils.rs | 53 +++++- .../test_files/array/array_append.slt | 62 +++++++ .../test_files/array/array_concat.slt | 19 +++ .../test_files/array/array_prepend.slt | 62 +++++++ .../test_files/array/array_replace.slt | 87 ++++++++++ 8 files changed, 505 insertions(+), 86 deletions(-) diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index 5dc437b3c20b5..1f03a0b17014e 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -20,26 +20,29 @@ use std::sync::Arc; use crate::make_array::make_array_inner; -use crate::utils::{align_array_dimensions, check_datatypes, make_scalar_function}; +use crate::utils::{ + align_array_dimensions, check_datatypes, list_inner_field, list_type_with_element, + make_scalar_function, +}; use arrow::array::{ Array, ArrayData, ArrayRef, Capacities, GenericListArray, MutableArrayData, OffsetSizeTrait, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::datatypes::{DataType, Field}; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::Result; use datafusion_common::utils::{ ListCoercion, base_type, coerced_type_with_base_type_only, }; use datafusion_common::{ cast::as_generic_list_array, - exec_err, plan_err, + exec_err, internal_err, plan_err, utils::{list_ndims, take_function_args}, }; use datafusion_expr::binary::type_union_resolution; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, Volatility, }; use datafusion_macros::user_doc; use itertools::Itertools; @@ -104,17 +107,26 @@ impl ScalarUDFImpl for ArrayAppend { &self.signature } - fn return_type(&self, arg_types: &[DataType]) -> Result { - let [array_type, element_type] = take_function_args(self.name(), arg_types)?; - if array_type.is_null() { - Ok(DataType::new_list(element_type.clone(), true)) - } else { - Ok(array_type.clone()) - } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let [array_field, element_field] = + take_function_args(self.name(), args.arg_fields)?; + let data_type = append_prepend_return_type( + array_field.data_type(), + element_field.data_type(), + element_field.is_nullable(), + ); + Ok(Arc::new(Field::new(self.name(), data_type, true))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_append_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + make_scalar_function(|args: &[ArrayRef]| array_append_inner(args, &return_type))( + &args.args, + ) } fn aliases(&self) -> &[String] { @@ -186,17 +198,26 @@ impl ScalarUDFImpl for ArrayPrepend { &self.signature } - fn return_type(&self, arg_types: &[DataType]) -> Result { - let [element_type, array_type] = take_function_args(self.name(), arg_types)?; - if array_type.is_null() { - Ok(DataType::new_list(element_type.clone(), true)) - } else { - Ok(array_type.clone()) - } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let [element_field, array_field] = + take_function_args(self.name(), args.arg_fields)?; + let data_type = append_prepend_return_type( + array_field.data_type(), + element_field.data_type(), + element_field.is_nullable(), + ); + Ok(Arc::new(Field::new(self.name(), data_type, true))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_prepend_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + make_scalar_function(|args: &[ArrayRef]| array_prepend_inner(args, &return_type))( + &args.args, + ) } fn aliases(&self) -> &[String] { @@ -375,13 +396,38 @@ pub fn array_concat_inner(args: &[ArrayRef]) -> Result { args[0].len(), ))) } else if large_list { - concat_internal::(args) + concat_internal::(args, None) + } else { + concat_internal::(args, None) + } +} + +/// Return type shared by `array_append` and `array_prepend`: the input list +/// type, except that its inner field is nullable whenever the appended or +/// prepended element may be null. +fn append_prepend_return_type( + array_type: &DataType, + element_type: &DataType, + element_nullable: bool, +) -> DataType { + if array_type.is_null() { + DataType::new_list(element_type.clone(), true) } else { - concat_internal::(args) + list_type_with_element(array_type, element_nullable) } } -fn concat_internal(args: &[ArrayRef]) -> Result { +/// 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. +fn concat_internal( + args: &[ArrayRef], + field: Option<&FieldRef>, +) -> Result { let args = align_array_dimensions::(args.to_vec())?; let list_arrays = args @@ -438,11 +484,14 @@ fn concat_internal(args: &[ArrayRef]) -> Result { offsets.push(O::usize_as(mutable.len())); } - let data_type = list_arrays[0].value_type(); + let field = match field { + Some(field) => Arc::clone(field), + None => Arc::new(Field::new_list_field(list_arrays[0].value_type(), true)), + }; let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(data_type, true)), + field, OffsetBuffer::new(offsets.into()), arrow::array::make_array(data), valid, @@ -451,22 +500,26 @@ fn concat_internal(args: &[ArrayRef]) -> Result { // Kernel functions -fn array_append_inner(args: &[ArrayRef]) -> Result { +fn array_append_inner(args: &[ArrayRef], return_type: &DataType) -> Result { let [array, values] = take_function_args("array_append", args)?; match array.data_type() { DataType::Null => make_array_inner(&[Arc::clone(values)]), - DataType::List(_) => general_append_and_prepend::(args, true), - DataType::LargeList(_) => general_append_and_prepend::(args, true), + DataType::List(_) => general_append_and_prepend::(args, true, return_type), + DataType::LargeList(_) => { + general_append_and_prepend::(args, true, return_type) + } arg_type => exec_err!("array_append does not support type {arg_type}"), } } -fn array_prepend_inner(args: &[ArrayRef]) -> Result { +fn array_prepend_inner(args: &[ArrayRef], return_type: &DataType) -> Result { let [values, array] = take_function_args("array_prepend", args)?; match array.data_type() { DataType::Null => make_array_inner(&[Arc::clone(values)]), - DataType::List(_) => general_append_and_prepend::(args, false), - DataType::LargeList(_) => general_append_and_prepend::(args, false), + DataType::List(_) => general_append_and_prepend::(args, false, return_type), + DataType::LargeList(_) => { + general_append_and_prepend::(args, false, return_type) + } arg_type => exec_err!("array_prepend does not support type {arg_type}"), } } @@ -474,6 +527,7 @@ fn array_prepend_inner(args: &[ArrayRef]) -> Result { fn general_append_and_prepend( args: &[ArrayRef], is_append: bool, + return_type: &DataType, ) -> Result where i64: TryInto, @@ -490,14 +544,22 @@ where (list_array, element_array) }; + let name = if is_append { + "array_append" + } else { + "array_prepend" + }; + let field = list_inner_field(name, return_type)?; + let res = match list_array.value_type() { - DataType::List(_) => concat_internal::(args)?, - DataType::LargeList(_) => concat_internal::(args)?, - data_type => { + DataType::List(_) | DataType::LargeList(_) => { + concat_internal::(args, Some(&field))? + } + _ => { return generic_append_and_prepend::( list_array, element_array, - &data_type, + field, is_append, ); } @@ -516,7 +578,7 @@ where /// /// * `list_array` - A reference to the ListArray to which elements will be appended/prepended. /// * `element_array` - A reference to the Array containing elements to be appended/prepended. -/// * `field` - A reference to the Field describing the data type of the arrays. +/// * `field` - The list field the output must carry, taken from the promised return type. /// * `is_append` - A boolean flag indicating whether to append (`true`) or prepend (`false`) elements. /// /// # Examples @@ -528,7 +590,7 @@ where fn generic_append_and_prepend( list_array: &GenericListArray, element_array: &ArrayRef, - data_type: &DataType, + field: FieldRef, is_append: bool, ) -> Result where @@ -565,7 +627,7 @@ where let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(data_type.to_owned(), true)), + field, OffsetBuffer::new(offsets.into()), arrow::array::make_array(data), None, diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs index 8f2ea1f40dcb2..cb7a316b289a9 100644 --- a/datafusion/functions-nested/src/extract.rs +++ b/datafusion/functions-nested/src/extract.rs @@ -47,7 +47,7 @@ use datafusion_expr::{ use datafusion_macros::user_doc; use std::sync::Arc; -use crate::utils::make_scalar_function; +use crate::utils::{list_inner_field, make_scalar_function}; // Create static instances of ScalarUDFs for each function make_udf_expr_and_func!( @@ -624,14 +624,7 @@ where // Carry the input's list field through to the output so that the returned // type matches the one promised by `return_type` / `return_field_from_args`, // including the field name, nullability and metadata. - let field = match array.data_type() { - List(field) | LargeList(field) => Arc::clone(field), - other => { - return internal_err!( - "general_array_slice got unexpected data type: {other}" - ); - } - }; + let field = list_inner_field("general_array_slice", array.data_type())?; // `use_nulls` is false because we never call `try_extend_nulls`: null rows are // emitted as empty slices. Arrow still allocates a validity buffer on its own diff --git a/datafusion/functions-nested/src/replace.rs b/datafusion/functions-nested/src/replace.rs index 71d6f578158f4..4bfd0c0dbecfe 100644 --- a/datafusion/functions-nested/src/replace.rs +++ b/datafusion/functions-nested/src/replace.rs @@ -22,17 +22,20 @@ use arrow::array::{ NullBufferBuilder, OffsetSizeTrait, Scalar, new_null_array, }; use arrow::buffer::OffsetBuffer; -use arrow::datatypes::{DataType, Field}; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cast::as_int64_array; use datafusion_common::utils::ListCoercion; -use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args}; +use datafusion_common::{ + Result, ScalarValue, exec_err, internal_err, utils::take_function_args, +}; use datafusion_expr::{ ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, - ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, + ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, }; use datafusion_macros::user_doc; -use crate::utils::compare_element_to_list; +use crate::utils::{compare_element_to_list, list_inner_field, list_type_with_element}; use std::sync::Arc; @@ -118,21 +121,28 @@ impl ScalarUDFImpl for ArrayReplace { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let return_type = args.return_field.data_type().clone(); let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?; let num_rows = args.number_rows; let list_array = list_arg.to_array(num_rows)?; match (from_arg, to_arg) { (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => { let result = array_replace_with_scalar_args( + self.name(), &list_array, scalar_from, scalar_to, 1i64, + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -140,10 +150,12 @@ impl ScalarUDFImpl for ArrayReplace { let from_array = from_arg.to_array(num_rows)?; let to_array = to_arg.to_array(num_rows)?; let result = array_replace_internal( + self.name(), &list_array, &from_array, &to_array, &[Some(1)], + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -217,11 +229,16 @@ impl ScalarUDFImpl for ArrayReplaceN { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let return_type = args.return_field.data_type().clone(); let [list_arg, from_arg, to_arg, max_arg] = take_function_args(self.name(), &args.args)?; let num_rows = args.number_rows; @@ -234,15 +251,17 @@ impl ScalarUDFImpl for ArrayReplaceN { ) => { let ScalarValue::Int64(Some(n)) = scalar_max else { return Ok(ColumnarValue::Array(new_null_array( - list_array.data_type(), + &return_type, num_rows, ))); }; let result = array_replace_with_scalar_args( + self.name(), &list_array, scalar_from, scalar_to, *n, + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -251,10 +270,12 @@ impl ScalarUDFImpl for ArrayReplaceN { let to_array = to_arg.to_array(num_rows)?; let max_array = max_arg.to_array(num_rows)?; let result = array_replace_n_inner( + self.name(), &list_array, &from_array, &to_array, &max_array, + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -326,21 +347,28 @@ impl ScalarUDFImpl for ArrayReplaceAll { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let return_type = args.return_field.data_type().clone(); let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?; let num_rows = args.number_rows; let list_array = list_arg.to_array(num_rows)?; match (from_arg, to_arg) { (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => { let result = array_replace_with_scalar_args( + self.name(), &list_array, scalar_from, scalar_to, i64::MAX, + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -348,10 +376,12 @@ impl ScalarUDFImpl for ArrayReplaceAll { let from_array = from_arg.to_array(num_rows)?; let to_array = to_arg.to_array(num_rows)?; let result = array_replace_internal( + self.name(), &list_array, &from_array, &to_array, &[Some(i64::MAX)], + &return_type, )?; Ok(ColumnarValue::Array(result)) } @@ -367,6 +397,24 @@ impl ScalarUDFImpl for ArrayReplaceAll { } } +/// Return field shared by `array_replace`, `array_replace_n` and +/// `array_replace_all`: the input list type, except that its inner field is +/// nullable whenever the replacement element may be null. +fn replace_return_field(name: &str, arg_fields: &[FieldRef]) -> Result { + // `array` is at index 0 and `to` at index 2 for all three functions. + // `from` never contributes values to the output, so `to` is the only + // argument besides `array` that can affect the output's type. + let [array_field, _from_field, to_field, ..] = arg_fields else { + return exec_err!( + "{name} expects at least 3 arguments, got {}", + arg_fields.len() + ); + }; + let data_type = + list_type_with_element(array_field.data_type(), to_field.is_nullable()); + Ok(Arc::new(Field::new(name, data_type, true))) +} + /// For each element of `list_array[i]`, replaces up to `arr_n[i]` occurrences /// of `from_array[i]`, `to_array[i]`. /// @@ -389,6 +437,7 @@ fn general_replace( from_array: &ArrayRef, to_array: &ArrayRef, arr_n: &[Option], + field: FieldRef, ) -> Result { // Build up the offsets for the final output array let mut offsets: Vec = Vec::with_capacity(list_array.len() + 1); @@ -502,7 +551,7 @@ fn general_replace( let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(list_array.value_type(), true)), + field, OffsetBuffer::::new(offsets.into()), arrow::array::make_array(data), valid.finish(), @@ -520,10 +569,17 @@ fn general_replace_with_scalar( needle: &Scalar, scalar_to: &ScalarValue, max_replacements: i64, + field: FieldRef, ) -> Result { - // No replacement needed - return unchanged. + // No replacement needed, but the output still has to carry the promised + // field, which may be more nullable than the input's. if max_replacements <= 0 { - return Ok(Arc::new(list_array.clone())); + return Ok(Arc::new(GenericListArray::::try_new( + field, + list_array.offsets().clone(), + Arc::clone(list_array.values()), + list_array.nulls().cloned(), + )?)); } let first_offset = list_array.offsets()[0].to_usize().unwrap(); @@ -598,7 +654,7 @@ fn general_replace_with_scalar( let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(list_array.value_type(), true)), + field, OffsetBuffer::new(offsets.into()), arrow::array::make_array(data), list_array.nulls().cloned(), @@ -609,10 +665,12 @@ fn general_replace_with_scalar( /// /// Uses a single bulk `not_distinct` comparison instead of per-row comparisons. fn array_replace_with_scalar_args( + name: &str, list_array: &ArrayRef, scalar_from: &ScalarValue, scalar_to: &ScalarValue, max_replacements: i64, + return_type: &DataType, ) -> Result { // `not_distinct` doesn't support nested types, fall back to the generic array path. if scalar_from.data_type().is_nested() { @@ -620,56 +678,74 @@ fn array_replace_with_scalar_args( let from_array = scalar_from.to_array_of_size(num_rows)?; let to_array = scalar_to.to_array_of_size(num_rows)?; return array_replace_internal( + name, list_array, &from_array, &to_array, &vec![Some(max_replacements); num_rows], + return_type, ); } let needle = Scalar::new(scalar_from.to_array_of_size(1)?); match list_array.data_type() { - DataType::List(_) => { - let list = list_array.as_list::(); - general_replace_with_scalar::(list, &needle, scalar_to, max_replacements) - } - DataType::LargeList(_) => { - let list = list_array.as_list::(); - general_replace_with_scalar::(list, &needle, scalar_to, max_replacements) - } - DataType::Null => Ok(new_null_array(list_array.data_type(), list_array.len())), - array_type => exec_err!("array_replace does not support type '{array_type}'."), + DataType::List(_) => general_replace_with_scalar::( + list_array.as_list::(), + &needle, + scalar_to, + max_replacements, + list_inner_field(name, return_type)?, + ), + DataType::LargeList(_) => general_replace_with_scalar::( + list_array.as_list::(), + &needle, + scalar_to, + max_replacements, + list_inner_field(name, return_type)?, + ), + DataType::Null => Ok(new_null_array(return_type, list_array.len())), + array_type => exec_err!("{name} does not support type '{array_type}'."), } } fn array_replace_internal( + name: &str, array: &ArrayRef, from: &ArrayRef, to: &ArrayRef, arr_n: &[Option], + return_type: &DataType, ) -> Result { match array.data_type() { - DataType::List(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, arr_n) - } - DataType::LargeList(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, arr_n) - } - DataType::Null => Ok(new_null_array(array.data_type(), array.len())), - array_type => exec_err!("array_replace does not support type '{array_type}'."), + DataType::List(_) => general_replace::( + array.as_list::(), + from, + to, + arr_n, + list_inner_field(name, return_type)?, + ), + DataType::LargeList(_) => general_replace::( + array.as_list::(), + from, + to, + arr_n, + list_inner_field(name, return_type)?, + ), + DataType::Null => Ok(new_null_array(return_type, array.len())), + array_type => exec_err!("{name} does not support type '{array_type}'."), } } fn array_replace_n_inner( + name: &str, array: &ArrayRef, from: &ArrayRef, to: &ArrayRef, max: &ArrayRef, + return_type: &DataType, ) -> Result { let arr_n = as_int64_array(max)?.iter().collect::>(); - array_replace_internal(array, from, to, &arr_n) + array_replace_internal(name, array, from, to, &arr_n, return_type) } #[cfg(test)] @@ -696,7 +772,14 @@ mod tests { Some(NullBuffer::from(vec![true, false])), )); - let result = array_replace_n_inner(&array, &from, &to, &max)?; + let result = array_replace_n_inner( + "array_replace_n", + &array, + &from, + &to, + &max, + array.data_type(), + )?; let expected = ListArray::from_iter_primitive::(vec![ Some(vec![Some(1), Some(9), Some(3)]), None, diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index 8b413686abcab..9822b6121e695 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -19,7 +19,7 @@ use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Fields}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; use arrow::array::{ Array, ArrayRef, BooleanArray, Float64Array, GenericListArray, NullBufferBuilder, @@ -35,6 +35,57 @@ use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err}; use datafusion_expr::ColumnarValue; use itertools::Itertools as _; +/// Computes the return type of a function that produces a list with the same +/// inner field as `array_type`, plus an element that may be null when +/// `element_nullable` is set. +/// +/// The inner field is carried over from `array_type` verbatim — name, metadata +/// and all — so that the type promised at planning time is the one the kernel +/// can actually build. Its nullability is widened when `element_nullable` is +/// set, because a nullable new element may introduce nulls into a list whose +/// elements were previously declared non-nullable. +/// +/// Types other than `List`/`LargeList` are returned unchanged; callers handle +/// `Null` themselves and the kernels reject anything else at execution time. +pub(crate) fn list_type_with_element( + array_type: &DataType, + element_nullable: bool, +) -> DataType { + match array_type { + DataType::List(field) => { + DataType::List(widen_nullability(field, element_nullable)) + } + DataType::LargeList(field) => { + DataType::LargeList(widen_nullability(field, element_nullable)) + } + other => other.clone(), + } +} + +fn widen_nullability(field: &FieldRef, nullable: bool) -> FieldRef { + if nullable && !field.is_nullable() { + Arc::new(field.as_ref().clone().with_nullable(true)) + } else { + Arc::clone(field) + } +} + +/// Extracts the inner field of a `List`/`LargeList` type, so that a kernel can +/// build a list array carrying exactly that field. +/// +/// Used both on an input's type and on the type promised by +/// [`ScalarUDFImpl::return_field_from_args`]. Anything else is a bug in the +/// caller's dispatch, hence the internal error; `context` names the kernel so +/// that error identifies where the bad dispatch happened. +/// +/// [`ScalarUDFImpl::return_field_from_args`]: datafusion_expr::ScalarUDFImpl::return_field_from_args +pub(crate) fn list_inner_field(context: &str, data_type: &DataType) -> Result { + match data_type { + DataType::List(field) | DataType::LargeList(field) => Ok(Arc::clone(field)), + other => internal_err!("{context} got unexpected data type: {other}"), + } +} + pub(crate) fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> { let data_type = args[0].data_type(); if !args.iter().all(|arg| { diff --git a/datafusion/sqllogictest/test_files/array/array_append.slt b/datafusion/sqllogictest/test_files/array/array_append.slt index 50949948c890e..0758a09a4925b 100644 --- a/datafusion/sqllogictest/test_files/array/array_append.slt +++ b/datafusion/sqllogictest/test_files/array/array_append.slt @@ -269,5 +269,67 @@ select array_append(column1, arrow_cast(make_array(1, 11, 111), 'FixedSizeList(3 [[1, 2, 3], [2, 9, 1], [7, 8, 9], [1, 2, 3], [1, 7, 4], [4, 5, 6], [1, 11, 111]] [[1, 2, 3], [11, 12, 13], [7, 8, 9]] [[4, 5, 6], [10, 11, 12], [4, 9, 8], [7, 8, 9], [10, 11, 12], [1, 8, 7], [1, 11, 111]] [[1, 2, 3], [11, 12, 13], [10, 11, 12]] +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ?T +select + array_append(arrow_cast(column1, 'List(Int64, field: ''element'')'), 4), + arrow_typeof(array_append(arrow_cast(column1, 'List(Int64, field: ''element'')'), 4)) +from values (make_array(1, 2, 3)); +---- +[1, 2, 3, 4] List(Int64, field: 'element') + +query ?T +select + array_append(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 4), + arrow_typeof(array_append(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 4)) +from values (make_array(1, 2, 3)); +---- +[1, 2, 3, 4] LargeList(Int64, field: 'element') + +# nested value types go through a different kernel path +query ?T +select + array_append(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2)), + arrow_typeof(array_append(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2))) +from values (make_array(make_array(1))); +---- +[[1], [2]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the appended element cannot be null +query ??TT +select + array_append(column1, 4), + array_append(arrow_cast(column1, 'LargeList(non-null Int64)'), 4), + arrow_typeof(array_append(column1, 4)), + arrow_typeof(array_append(arrow_cast(column1, 'LargeList(non-null Int64)'), 4)) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 3), 'List(non-null Int64)')) +; +---- +[4] [4] List(non-null Int64) LargeList(non-null Int64) +[4] [4] List(non-null Int64) LargeList(non-null Int64) +[1, 2, 3, 4] [1, 2, 3, 4] List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the appended element is nullable, since the result +# genuinely contains a null element +query ?T +select + array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), NULL), + arrow_typeof(array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), NULL)); +---- +[1, 2, NULL] List(Int64) + +query ?T +select + array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), column1), + arrow_typeof(array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), column1)) +from values (3), (NULL); +---- +[1, 2, 3] List(Int64) +[1, 2, NULL] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_concat.slt b/datafusion/sqllogictest/test_files/array/array_concat.slt index 168b307a1e636..5b7985ef6d194 100644 --- a/datafusion/sqllogictest/test_files/array/array_concat.slt +++ b/datafusion/sqllogictest/test_files/array/array_concat.slt @@ -419,5 +419,24 @@ 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 +query ?T +select + array_concat(arrow_cast(column1, 'List(non-null Int64)'), make_array(3)), + arrow_typeof(array_concat(arrow_cast(column1, 'List(non-null Int64)'), make_array(3))) +from values (make_array(1, 2)); +---- +[1, 2, 3] List(Int64) + +query ?T +select + array_concat(arrow_cast(column1, 'List(Int64, field: ''element'')'), make_array(3)), + arrow_typeof(array_concat(arrow_cast(column1, 'List(Int64, field: ''element'')'), make_array(3))) +from values (make_array(1, 2)); +---- +[1, 2, 3] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_prepend.slt b/datafusion/sqllogictest/test_files/array/array_prepend.slt index 14b53e93b3d0d..bfb61ab4f9f91 100644 --- a/datafusion/sqllogictest/test_files/array/array_prepend.slt +++ b/datafusion/sqllogictest/test_files/array/array_prepend.slt @@ -273,5 +273,67 @@ select array_prepend(arrow_cast(make_array(1, 11, 111), 'FixedSizeList(3, Int64) [[1, 11, 111], [1, 2, 3], [2, 9, 1], [7, 8, 9], [1, 2, 3], [1, 7, 4], [4, 5, 6]] [[7, 8, 9], [1, 2, 3], [11, 12, 13]] [[1, 11, 111], [4, 5, 6], [10, 11, 12], [4, 9, 8], [7, 8, 9], [10, 11, 12], [1, 8, 7]] [[10, 11, 12], [1, 2, 3], [11, 12, 13]] +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ?T +select + array_prepend(0, arrow_cast(column1, 'List(Int64, field: ''element'')')), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'List(Int64, field: ''element'')'))) +from values (make_array(1, 2, 3)); +---- +[0, 1, 2, 3] List(Int64, field: 'element') + +query ?T +select + array_prepend(0, arrow_cast(column1, 'LargeList(Int64, field: ''element'')')), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'LargeList(Int64, field: ''element'')'))) +from values (make_array(1, 2, 3)); +---- +[0, 1, 2, 3] LargeList(Int64, field: 'element') + +# nested value types go through a different kernel path +query ?T +select + array_prepend(make_array(0), arrow_cast(column1, 'List(List(Int64), field: ''element'')')), + arrow_typeof(array_prepend(make_array(0), arrow_cast(column1, 'List(List(Int64), field: ''element'')'))) +from values (make_array(make_array(1))); +---- +[[0], [1]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the prepended element cannot be null +query ??TT +select + array_prepend(0, column1), + array_prepend(0, arrow_cast(column1, 'LargeList(non-null Int64)')), + arrow_typeof(array_prepend(0, column1)), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'LargeList(non-null Int64)'))) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 3), 'List(non-null Int64)')) +; +---- +[0] [0] List(non-null Int64) LargeList(non-null Int64) +[0] [0] List(non-null Int64) LargeList(non-null Int64) +[0, 1, 2, 3] [0, 1, 2, 3] List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the prepended element is nullable, since the result +# genuinely contains a null element +query ?T +select + array_prepend(NULL, arrow_cast(make_array(1, 2), 'List(non-null Int64)')), + arrow_typeof(array_prepend(NULL, arrow_cast(make_array(1, 2), 'List(non-null Int64)'))); +---- +[NULL, 1, 2] List(Int64) + +query ?T +select + array_prepend(column1, arrow_cast(make_array(1, 2), 'List(non-null Int64)')), + arrow_typeof(array_prepend(column1, arrow_cast(make_array(1, 2), 'List(non-null Int64)'))) +from values (0), (NULL); +---- +[0, 1, 2] List(Int64) +[NULL, 1, 2] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_replace.slt b/datafusion/sqllogictest/test_files/array/array_replace.slt index cab84007bcd53..ce45e6440dddf 100644 --- a/datafusion/sqllogictest/test_files/array/array_replace.slt +++ b/datafusion/sqllogictest/test_files/array/array_replace.slt @@ -753,6 +753,93 @@ select ---- [3, 5, NULL] [3, 5, 5] [3, 5, 5] +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ???TTT +select + array_replace(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9), + array_replace_n(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9, 1), + array_replace_all(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9), + arrow_typeof(array_replace(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9)) +from values (make_array(1, 2, 2)); +---- +[1, 9, 2] [1, 9, 2] [1, 9, 9] List(Int64, field: 'element') List(Int64, field: 'element') List(Int64, field: 'element') + +query ???TTT +select + array_replace(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9), + array_replace_n(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9, 1), + array_replace_all(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9), + arrow_typeof(array_replace(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9)) +from values (make_array(1, 2, 2)); +---- +[1, 9, 2] [1, 9, 2] [1, 9, 9] LargeList(Int64, field: 'element') LargeList(Int64, field: 'element') LargeList(Int64, field: 'element') + +# nested from/to values fall back to the generic comparison path +query ?T +select + array_replace(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2), make_array(9)), + arrow_typeof(array_replace(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2), make_array(9))) +from values (make_array(make_array(1), make_array(2))); +---- +[[1], [9]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the replacement cannot be null +query ???TTT +select + array_replace(column1, 2, 9), + array_replace_n(column1, 2, 9, 1), + array_replace_all(arrow_cast(column1, 'LargeList(non-null Int64)'), 2, 9), + arrow_typeof(array_replace(column1, 2, 9)), + arrow_typeof(array_replace_n(column1, 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'LargeList(non-null Int64)'), 2, 9)) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 2), 'List(non-null Int64)')) +; +---- +[] [] [] List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) +NULL NULL NULL List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) +[1, 9, 2] [1, 9, 2] [1, 9, 9] List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the replacement is nullable, since the result +# genuinely contains a null element +query ???TTT +select + array_replace(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL), + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 1), + array_replace_all(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL), + arrow_typeof(array_replace(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL)) +from values (make_array(1, 2, 2)); +---- +[1, NULL, 2] [1, NULL, 2] [1, NULL, NULL] List(Int64) List(Int64) List(Int64) + +# a max of 0 short circuits without replacing anything, but must still return +# the promised (widened) type +query ?T +select + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 0), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 0)) +from values (make_array(1, 2, 2)); +---- +[1, 2, 2] List(Int64) + +# a NULL max yields a NULL row of the promised type +query ?T +select + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, 9, NULL), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, 9, NULL)) +from values (make_array(1, 2, 2)); +---- +NULL List(non-null Int64) + statement ok