diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index c3dc4c67cf12c..8d06140889a55 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -317,10 +317,23 @@ impl ScalarUDFImpl for ArrayConcat { } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - let base_type = base_type(&self.return_type(arg_types)?); + 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| { - coerced_type_with_base_type_only(arg_type, &base_type, coercion) + 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, + } }); Ok(arg_types.collect()) diff --git a/datafusion/sqllogictest/test_files/array/array_concat.slt b/datafusion/sqllogictest/test_files/array/array_concat.slt index 0f847811615c7..168b307a1e636 100644 --- a/datafusion/sqllogictest/test_files/array/array_concat.slt +++ b/datafusion/sqllogictest/test_files/array/array_concat.slt @@ -121,6 +121,38 @@ select ---- [1, 2, 3] List(Utf8View) +# Concatenating mixed list and large list — return type widens to LargeList +query ?T +select + array_concat(make_array(1, 2), arrow_cast([3, 4], 'LargeList(Int64)')), + arrow_typeof(array_concat(make_array(1, 2), arrow_cast([3, 4], 'LargeList(Int64)'))); +---- +[1, 2, 3, 4] LargeList(Int64) + +# Reverse argument order: LargeList first, plain list second +query ?T +select + array_concat(arrow_cast([1, 2], 'LargeList(Int64)'), make_array(3, 4)), + arrow_typeof(array_concat(arrow_cast([1, 2], 'LargeList(Int64)'), make_array(3, 4))); +---- +[1, 2, 3, 4] LargeList(Int64) + +# FixedSizeList combined with LargeList — also widens to LargeList +query ?T +select + array_concat(arrow_cast([1, 2], 'FixedSizeList(2, Int64)'), arrow_cast([3, 4], 'LargeList(Int64)')), + arrow_typeof(array_concat(arrow_cast([1, 2], 'FixedSizeList(2, Int64)'), arrow_cast([3, 4], 'LargeList(Int64)'))); +---- +[1, 2, 3, 4] LargeList(Int64) + +# Three-way mix: List, LargeList, List +query ?T +select + array_concat(make_array(1, 2), arrow_cast([3], 'LargeList(Int64)'), make_array(4, 5)), + arrow_typeof(array_concat(make_array(1, 2), arrow_cast([3], 'LargeList(Int64)'), make_array(4, 5))); +---- +[1, 2, 3, 4, 5] LargeList(Int64) + # array_concat with NULL elements inside arrays query ? select array_concat([1, NULL, 3], [NULL, 5]);