Fix array_concat with NULL arrays (apache#16348) (#32)

Author: gabotechs

datafusion/functions-nested/src/concat.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 use crate::make_array::make_array_inner;
 use crate::utils::{align_array_dimensions, check_datatypes, make_scalar_function};
 use arrow::array::{
-    Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, NullArray,
+    Array, ArrayData, ArrayRef, Capacities, GenericListArray, MutableArrayData,
     NullBufferBuilder, OffsetSizeTrait,
 };
 use arrow::buffer::OffsetBuffer;
@@ -42,6 +42,7 @@ use datafusion_expr::{
     ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
+use itertools::Itertools;
 
 make_udf_expr_and_func!(
     ArrayAppend,
@@ -365,11 +366,23 @@ pub(crate) fn array_concat_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
             _ => (),
         }
 
-        all_null = false
+        if arg.null_count() < arg.len() {
+            all_null = false;
+        }
     }
 
     if all_null {
-        Ok(Arc::new(NullArray::new(args[0].len())))
+        // 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::<i64>(args)
     } else {

datafusion/sqllogictest/test_files/array.slt

@@ -3070,6 +3070,42 @@ select array_concat([]);
 ----
 []
 
+# test with NULL array
+query ?
+select array_concat(NULL::integer[]);
+----
+NULL
+
+# test with multiple NULL arrays
+query ?
+select array_concat(NULL::integer[], NULL::integer[]);
+----
+NULL
+
+# test with NULL LargeList
+query ?
+select array_concat(arrow_cast(NULL::string[], 'LargeList(Utf8)'));
+----
+NULL
+
+# test with NULL FixedSizeList
+query ?
+select array_concat(arrow_cast(NULL::string[], 'FixedSizeList(2, Utf8)'));
+----
+NULL
+
+# test with mix of NULL and empty arrays
+query ?
+select array_concat(NULL::integer[], []);
+----
+[]
+
+# test with mix of NULL and non-empty arrays
+query ?
+select array_concat(NULL::integer[], [1, 2, 3]);
+----
+[1, 2, 3]
+
 # Concatenating strings arrays
 query ?
 select array_concat(