[AURON #2343] Support native flatten for mixed child array nullability (#2344)

# Which issue does this PR close?

Closes #2343 

# Rationale for this change

Auron currently excludes Spark's `flatten function` tests because native
execution can fail when child arrays have different `containsNull`
metadata.

`flatten` is a built-in Spark array expression, so supporting it in
native execution improves Spark SQL compatibility and reduces fallback
coverage gaps.

# What changes are included in this PR?

- Adds a native `Spark_ArrayFlatten` function.
- Converts Spark `Flatten` expressions to the new native function.
- Preserves Spark semantics for:
  - null outer arrays
  - null child arrays
  - empty arrays
  - null elements inside child arrays
  - primitive and struct element arrays
- Adds an Auron SQL test for mixed child array nullability.
- Removes the `flatten function` excludes from Spark 3.1/3.2/3.4/3.5
test settings.

# Are there any user-facing changes?
No user-facing API changes. More `flatten` expressions can now be
executed natively by Auron.

# How was this patch tested?
UT.

Signed-off-by: weimingdiit <weimingdiit@gmail.com>
Co-authored-by: Shilun Fan <slfan1989@apache.org>

Author: weimingdiit

auron-spark-tests/spark31/src/test/scala/org/apache/auron/utils/AuronSparkTestSettings.scala

@@ -31,8 +31,6 @@ class AuronSparkTestSettings extends SparkTestSettings {
     // Native execution wraps SparkRuntimeException from map_concat input validation in
     // SparkException.
     .exclude("map_concat function")
-    // Native flatten can fail when child arrays have different containsNull metadata.
-    .exclude("flatten function")
     // Native execution wraps SparkRuntimeException from null map keys in SparkException.
     .exclude("SPARK-24734: Fix containsNull of Concat for array type")
 

auron-spark-tests/spark32/src/test/scala/org/apache/auron/utils/AuronSparkTestSettings.scala

@@ -31,8 +31,6 @@ class AuronSparkTestSettings extends SparkTestSettings {
     .exclude("map with arrays")
     .exclude("map_concat function")
     .exclude("SPARK-24734: Fix containsNull of Concat for array type")
-    // Native flatten can fail when child arrays have different containsNull metadata.
-    .exclude("flatten function")
 
   enableSuite[AuronDateFunctionsSuite]
     // Native execution wraps Spark parsing/format validation exceptions in SparkException.

auron-spark-tests/spark34/src/test/scala/org/apache/auron/utils/AuronSparkTestSettings.scala

@@ -32,8 +32,6 @@ class AuronSparkTestSettings extends SparkTestSettings {
     .exclude("SPARK-24734: Fix containsNull of Concat for array type")
     .exclude("array_insert functions")
     .exclude("transform keys function - Invalid lambda functions and exceptions")
-    // Native flatten can fail when child arrays have different containsNull metadata.
-    .exclude("flatten function")
 
   enableSuite[AuronDateFunctionsSuite]
     // Native execution wraps Spark parsing/format validation exceptions in SparkException.

auron-spark-tests/spark35/src/test/scala/org/apache/auron/utils/AuronSparkTestSettings.scala

@@ -32,8 +32,6 @@ class AuronSparkTestSettings extends SparkTestSettings {
     .exclude("SPARK-24734: Fix containsNull of Concat for array type")
     .exclude("array_insert functions")
     .exclude("transform keys function - Invalid lambda functions and exceptions")
-    // Native flatten can fail when child arrays have different containsNull metadata.
-    .exclude("flatten function")
 
   enableSuite[AuronDateFunctionsSuite]
     // Native execution wraps Spark parsing/format validation exceptions in SparkException.

native-engine/datafusion-ext-functions/src/lib.rs

@@ -65,6 +65,7 @@ pub fn create_auron_ext_function(
         }
         "Spark_ParseJson" => Arc::new(spark_get_json_object::spark_parse_json),
         "Spark_ArrayReverse" => Arc::new(spark_array::array_reverse),
+        "Spark_ArrayFlatten" => Arc::new(spark_array::array_flatten),
         "Spark_MakeArray" => Arc::new(spark_make_array::array),
         "Spark_MapConcat" => Arc::new(spark_map::map_concat),
         "Spark_MapFromArrays" => Arc::new(spark_map::map_from_arrays),

native-engine/datafusion-ext-functions/src/spark_array.rs

@@ -13,9 +13,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//! Array expressions
+
 use std::sync::Arc;
 
-use arrow::{array::*, datatypes::DataType};
+use arrow::{
+    array::{Array, ArrayRef, ListArray, MutableArrayData, make_array},
+    buffer::{NullBuffer, OffsetBuffer, ScalarBuffer},
+    datatypes::{DataType, Field},
+};
 use datafusion::{
     common::{Result, ScalarValue},
     logical_expr::ColumnarValue,
@@ -80,11 +86,94 @@ fn reverse_list_array(array: &ListArray) -> Result<ArrayRef> {
     ScalarValue::iter_to_array(values)
 }
 
+fn as_list_array(array: &ArrayRef) -> Result<ListArray> {
+    downcast_any!(array, ListArray).cloned()
+}
+
+fn columnar_value_to_list_array(arg: &ColumnarValue) -> Result<ListArray> {
+    match arg {
+        ColumnarValue::Array(array) if matches!(array.data_type(), DataType::Null) => {
+            Ok(ListArray::new_null(
+                Arc::new(Field::new_list_field(DataType::Null, true)),
+                array.len(),
+            ))
+        }
+        ColumnarValue::Array(array) => as_list_array(array),
+        ColumnarValue::Scalar(scalar) if scalar.is_null() => {
+            let list_field = match scalar.data_type() {
+                DataType::List(field) => field,
+                _ => Arc::new(Field::new_list_field(DataType::Null, true)),
+            };
+            Ok(ListArray::new_null(list_field, 1))
+        }
+        ColumnarValue::Scalar(scalar) => {
+            let array = scalar.to_array()?;
+            as_list_array(&array)
+        }
+    }
+}
+
+/// Flatten an array of arrays into a single array.
+pub fn array_flatten(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    if args.len() != 1 {
+        df_execution_err!("array_flatten expects one argument, got {}", args.len())?;
+    }
+
+    let outer = columnar_value_to_list_array(&args[0])?;
+    let inner = as_list_array(&outer.values())?;
+    let inner_values = inner.values();
+    let inner_values_data = inner_values.to_data();
+    let mut mutable = MutableArrayData::new(vec![&inner_values_data], true, inner_values.len());
+
+    let mut offsets = Vec::with_capacity(outer.len() + 1);
+    let mut valids = Vec::with_capacity(outer.len());
+    let mut offset = 0i32;
+    offsets.push(offset);
+
+    for row_idx in 0..outer.len() {
+        if outer.is_null(row_idx) {
+            valids.push(false);
+            offsets.push(offset);
+            continue;
+        }
+
+        let outer_start = outer.value_offsets()[row_idx] as usize;
+        let outer_end = outer.value_offsets()[row_idx + 1] as usize;
+        let row_valid = (outer_start..outer_end).all(|inner_idx| inner.is_valid(inner_idx));
+
+        valids.push(row_valid);
+        if row_valid {
+            let mut row_len = 0i32;
+            for inner_idx in outer_start..outer_end {
+                let inner_start = inner.value_offsets()[inner_idx] as usize;
+                let inner_end = inner.value_offsets()[inner_idx + 1] as usize;
+                mutable.extend(0, inner_start, inner_end);
+                row_len += (inner_end - inner_start) as i32;
+            }
+            offset += row_len;
+        }
+        offsets.push(offset);
+    }
+
+    let values = make_array(mutable.freeze());
+    let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
+    Ok(ColumnarValue::Array(Arc::new(ListArray::try_new(
+        field,
+        OffsetBuffer::new(ScalarBuffer::from(offsets)),
+        values,
+        Some(NullBuffer::from(valids)),
+    )?)))
+}
+
 #[cfg(test)]
 mod test {
-    use std::sync::Arc;
+    use std::{error::Error, sync::Arc};
 
-    use arrow::datatypes::Int32Type;
+    use arrow::{
+        array::{ArrayRef, ListArray},
+        buffer::{NullBuffer, OffsetBuffer, ScalarBuffer},
+        datatypes::{Field, Int32Type},
+    };
     use datafusion::common::ScalarValue;
 
     use super::*;
@@ -137,4 +226,38 @@ mod test {
         assert_eq!(&result, &expected);
         Ok(())
     }
+
+    #[test]
+    fn test_array_flatten_int() -> std::result::Result<(), Box<dyn Error>> {
+        let inner = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+            Some(vec![Some(1), None]),
+            Some(vec![Some(3)]),
+            Some(vec![]),
+            Some(vec![Some(4), Some(5)]),
+            Some(vec![Some(6)]),
+            None,
+        ]);
+        let inner: ArrayRef = Arc::new(inner);
+        let input: ArrayRef = Arc::new(ListArray::try_new(
+            Arc::new(Field::new_list_field(inner.data_type().clone(), true)),
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4, 6, 6, 6])),
+            inner,
+            Some(NullBuffer::from(vec![true, true, true, true, false])),
+        )?);
+
+        let result = array_flatten(&[ColumnarValue::Array(input)])?.into_array(5)?;
+
+        let expected = vec![
+            Some(vec![Some(1), None, Some(3)]),
+            Some(vec![Some(4), Some(5)]),
+            None,
+            Some(vec![]),
+            None,
+        ];
+        let expected: ArrayRef =
+            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(expected));
+
+        assert_eq!(&result, &expected);
+        Ok(())
+    }
 }

native-engine/datafusion-ext-functions/src/spark_make_array.rs

@@ -17,12 +17,15 @@
 
 use std::sync::Arc;
 
-use arrow::{array::*, datatypes::DataType};
+use arrow::{
+    array::*,
+    datatypes::{DataType, Field, Fields},
+};
 use datafusion::{
     common::{Result, ScalarValue},
     logical_expr::ColumnarValue,
 };
-use datafusion_ext_commons::df_execution_err;
+use datafusion_ext_commons::{arrow::cast::cast, df_execution_err};
 
 macro_rules! downcast_vec {
     ($ARGS:expr, $ARRAY_TYPE:ident) => {{
@@ -105,9 +108,20 @@ fn array_array(args: &[ArrayRef]) -> Result<ArrayRef> {
         DataType::UInt16 => array!(args, UInt16Array, UInt16Builder),
         DataType::UInt32 => array!(args, UInt32Array, UInt32Builder),
         DataType::UInt64 => array!(args, UInt64Array, UInt64Builder),
-        data_type => {
+        _ => {
             // naive implementation with scalar values
             let num_rows = args[0].len();
+            let data_type = common_array_element_data_type(args)?;
+            let args = args
+                .iter()
+                .map(|arg| {
+                    if arg.data_type() == &data_type {
+                        Ok(arg.clone())
+                    } else {
+                        cast(arg.as_ref(), &data_type)
+                    }
+                })
+                .collect::<Result<Vec<_>>>()?;
             let mut output_scalars = Vec::with_capacity(num_rows);
             for i in 0..num_rows {
                 let row_scalars: Vec<ScalarValue> = args
@@ -116,7 +130,7 @@ fn array_array(args: &[ArrayRef]) -> Result<ArrayRef> {
                     .collect::<Result<_>>()?;
                 output_scalars.push(ScalarValue::List(ScalarValue::new_list(
                     &row_scalars,
-                    data_type,
+                    &data_type,
                     true,
                 )));
             }
@@ -126,6 +140,57 @@ fn array_array(args: &[ArrayRef]) -> Result<ArrayRef> {
     Ok(res)
 }
 
+fn common_array_element_data_type(args: &[ArrayRef]) -> Result<DataType> {
+    let mut data_type = args[0].data_type().clone();
+    for arg in &args[1..] {
+        data_type = widen_data_type(&data_type, arg.data_type())?;
+    }
+    Ok(data_type)
+}
+
+fn widen_data_type(left: &DataType, right: &DataType) -> Result<DataType> {
+    if left == right {
+        return Ok(left.clone());
+    }
+
+    match (left, right) {
+        (DataType::List(left_field), DataType::List(right_field)) => {
+            let item_type = widen_data_type(left_field.data_type(), right_field.data_type())?;
+            Ok(DataType::List(Arc::new(Field::new(
+                left_field.name().as_str(),
+                item_type,
+                left_field.is_nullable() || right_field.is_nullable(),
+            ))))
+        }
+        (DataType::Struct(left_fields), DataType::Struct(right_fields))
+            if left_fields.len() == right_fields.len() =>
+        {
+            let fields = left_fields
+                .iter()
+                .zip(right_fields.iter())
+                .map(|(left_field, right_field)| {
+                    if left_field.name() != right_field.name() {
+                        return df_execution_err!(
+                            "array child struct fields must have same names, got {} and {}",
+                            left_field.name(),
+                            right_field.name()
+                        );
+                    }
+                    Ok(Arc::new(Field::new(
+                        left_field.name().as_str(),
+                        widen_data_type(left_field.data_type(), right_field.data_type())?,
+                        left_field.is_nullable() || right_field.is_nullable(),
+                    )))
+                })
+                .collect::<Result<Vec<_>>>()?;
+            Ok(DataType::Struct(Fields::from(fields)))
+        }
+        _ => df_execution_err!(
+            "array child values must have same data type, got {left:?} and {right:?}"
+        ),
+    }
+}
+
 /// put values in an array.
 pub fn array(values: &[ColumnarValue]) -> Result<ColumnarValue> {
     let arrays: Vec<ArrayRef> = values
@@ -145,7 +210,8 @@ mod test {
 
     use arrow::{
         array::{ArrayRef, Int32Array, ListArray},
-        datatypes::{Float32Type, Int32Type},
+        buffer::{OffsetBuffer, ScalarBuffer},
+        datatypes::{DataType, Field, Float32Type, Int32Type},
     };
     use datafusion::{common::ScalarValue, physical_plan::ColumnarValue};
 
@@ -215,4 +281,48 @@ mod test {
         assert_eq!(&result, &expected);
         Ok(())
     }
+
+    #[test]
+    fn test_make_array_with_mixed_child_array_nullability() -> Result<(), Box<dyn Error>> {
+        let non_nullable_items: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2)]));
+        let non_nullable_list: ArrayRef = Arc::new(ListArray::try_new(
+            Arc::new(Field::new_list_field(DataType::Int32, false)),
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2])),
+            non_nullable_items,
+            None,
+        )?);
+
+        let nullable_items: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(3)]));
+        let nullable_list: ArrayRef = Arc::new(ListArray::try_new(
+            Arc::new(Field::new_list_field(DataType::Int32, true)),
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2])),
+            nullable_items,
+            None,
+        )?);
+
+        let result = array(&vec![
+            ColumnarValue::Array(non_nullable_list),
+            ColumnarValue::Array(nullable_list),
+        ])?
+        .into_array(2)?;
+
+        let expected_values: ArrayRef = Arc::new(ListArray::try_new(
+            Arc::new(Field::new_list_field(DataType::Int32, true)),
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2, 3, 4])),
+            Arc::new(Int32Array::from(vec![Some(1), None, Some(2), Some(3)])),
+            None,
+        )?);
+        let expected: ArrayRef = Arc::new(ListArray::try_new(
+            Arc::new(Field::new_list_field(
+                expected_values.data_type().clone(),
+                true,
+            )),
+            OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 4])),
+            expected_values,
+            None,
+        )?);
+
+        assert_eq!(&result, &expected);
+        Ok(())
+    }
 }

spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala

@@ -96,6 +96,31 @@ class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite {
     }
   }
 
+  test("flatten function with mixed child array nullability") {
+    withTable("t1") {
+      sql("""
+          |create table t1 using parquet as
+          |select
+          |  array(array(1, 2), array(3, cast(null as int)), array()) as ints,
+          |  array(array(named_struct('a', 1)), array(named_struct('a', 2))) as structs
+          |union all
+          |select
+          |  array(array(4), cast(null as array<int>)),
+          |  array(array(named_struct('a', 3)), cast(array() as array<struct<a:int>>))
+          |union all
+          |select
+          |  cast(array() as array<array<int>>),
+          |  cast(array() as array<array<struct<a:int>>>)
+          |union all
+          |select
+          |  cast(null as array<array<int>>),
+          |  cast(null as array<array<struct<a:int>>>)
+          |""".stripMargin)
+
+      checkSparkAnswerAndOperator("select flatten(ints), flatten(structs) from t1")
+    }
+  }
+
   test("expm1 function") {
     withTable("t1") {
       sql("create table t1(c1 double) using parquet")