Feat: Add support for array_min, array_max, sort_array, array_zip & array_union

native/core/src/execution/planner.rs

@@ -84,10 +84,10 @@ use datafusion_comet_proto::{
     spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning},
 };
 use datafusion_comet_spark_expr::{
-    ArrayInsert, Avg, AvgDecimal, BitwiseNotExpr, Cast, CheckOverflow, Contains, Correlation,
-    Covariance, CreateNamedStruct, DateTruncExpr, EndsWith, GetArrayStructFields, GetStructField,
-    HourExpr, IfExpr, Like, ListExtract, MinuteExpr, NormalizeNaNAndZero, RLike, SecondExpr,
-    SparkCastOptions, StartsWith, Stddev, StringSpaceExpr, SubstringExpr, SumDecimal,
+    ArrayInsert, ArrayMinMax, Avg, AvgDecimal, BitwiseNotExpr, Cast, CheckOverflow, Contains,
+    Correlation, Covariance, CreateNamedStruct, DateTruncExpr, EndsWith, GetArrayStructFields,
+    GetStructField, HourExpr, IfExpr, Like, ListExtract, MinuteExpr, NormalizeNaNAndZero, RLike,
+    SecondExpr, SparkCastOptions, StartsWith, Stddev, StringSpaceExpr, SubstringExpr, SumDecimal,
     TimestampTruncExpr, ToJson, UnboundColumn, Variance,
 };
 use datafusion_common::scalar::ScalarStructBuilder;
@@ -100,6 +100,8 @@ use datafusion_expr::{
     WindowFunctionDefinition,
 };
 use datafusion_functions_nested::array_has::ArrayHas;
+use datafusion_functions_nested::set_ops::array_union_udf;
+use datafusion_functions_nested::sort::array_sort_udf;
 use datafusion_physical_expr::expressions::{Literal, StatsType};
 use datafusion_physical_expr::window::WindowExpr;
 use datafusion_physical_expr::LexOrdering;
@@ -735,6 +737,66 @@ impl PhysicalPlanner {
                 ));
                 Ok(array_has_expr)
             }
+            ExprStruct::ArrayMin(expr) => {
+                let src_array_expr =
+                    self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                Ok(Arc::new(ArrayMinMax::new(src_array_expr, false, true)))
+            }
+            ExprStruct::ArrayMax(expr) => {
+                let src_array_expr =
+                    self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                Ok(Arc::new(ArrayMinMax::new(src_array_expr, false, true)))
+            }
+            ExprStruct::SortArray(expr) => {
+                let array_expr =
+                    self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let asc_order_expr =
+                    self.create_expr(expr.right.as_ref().unwrap(), Arc::clone(&input_schema))?;
+
+                let true_literal_expr: Arc<dyn PhysicalExpr> =
+                    Arc::new(Literal::new(ScalarValue::Boolean(Some(true))));
+
+                let is_ordering_ignored: Arc<dyn PhysicalExpr> =
+                    Arc::new(IsNullExpr::new(Arc::clone(&asc_order_expr)));
+
+                // desc order expr
+                let not_asc_order_expr: Arc<dyn PhysicalExpr> =
+                    Arc::new(NotExpr::new(Arc::clone(&asc_order_expr)));
+
+                // Default to true of ordering expression is ignored
+                let ordering_expr: Arc<dyn PhysicalExpr> = Arc::new(CaseExpr::try_new(
+                    None,
+                    vec![(is_ordering_ignored, Arc::clone(&true_literal_expr))],
+                    Some(not_asc_order_expr),
+                )?);
+
+                let nulls_first: Arc<dyn PhysicalExpr> = true_literal_expr;
+                let return_type = array_expr.data_type(&input_schema)?;
+                let args = vec![array_expr, ordering_expr, nulls_first];
+                Ok(Arc::new(ScalarFunctionExpr::new(
+                    "array_sort",
+                    array_sort_udf(),
+                    args,
+                    return_type,
+                )))
+            }
+            ExprStruct::ArrayZip(expr) => {
+                unimplemented!()
+            }
+            ExprStruct::ArrayUnion(expr) => {
+                let left =
+                    self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let right =
+                    self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                let return_type = left.data_type(&input_schema)?;
+                let args = vec![left, right];
+                Ok(Arc::new(ScalarFunctionExpr::new(
+                    "array_union",
+                    array_union_udf(),
+                    args,
+                    return_type,
+                )))
+            }
             expr => Err(ExecutionError::GeneralError(format!(
                 "Not implemented: {:?}",
                 expr

native/proto/src/proto/expr.proto

@@ -85,6 +85,12 @@ message Expr {
     BinaryExpr array_append = 58;
     ArrayInsert array_insert = 59;
     BinaryExpr array_contains = 60;
+    UnaryExpr array_min = 61;
+    UnaryExpr array_max = 62;
+    BinaryExpr sort_array = 63;
+    BinaryExpr array_position = 64;
+    BinaryExpr array_zip = 65;
+    BinaryExpr array_union = 66;
   }
 }
 

native/spark-expr/src/lib.rs

@@ -72,7 +72,7 @@ pub use cast::{spark_cast, Cast, SparkCastOptions};
 pub use comet_scalar_funcs::create_comet_physical_fun;
 pub use error::{SparkError, SparkResult};
 pub use if_expr::IfExpr;
-pub use list::{ArrayInsert, GetArrayStructFields, ListExtract};
+pub use list::{ArrayInsert, ArrayMinMax, GetArrayStructFields, ListExtract};
 pub use regexp::RLike;
 pub use struct_funcs::*;
 

native/spark-expr/src/list.rs

@@ -18,12 +18,18 @@
 use arrow::{
     array::{as_primitive_array, Capacities, MutableArrayData},
     buffer::{NullBuffer, OffsetBuffer},
+    compute,
     datatypes::ArrowNativeType,
     record_batch::RecordBatch,
 };
+use arrow_array::builder::PrimitiveBuilder;
+use arrow_array::cast::AsArray;
+use arrow_array::types::Int32Type;
 use arrow_array::{
-    make_array, Array, ArrayRef, GenericListArray, Int32Array, OffsetSizeTrait, StructArray,
+    make_array, Array, ArrayRef, ArrowNumericType, ArrowPrimitiveType, GenericListArray,
+    Int32Array, OffsetSizeTrait, PrimitiveArray, StructArray,
 };
+use arrow_data::ArrayData;
 use arrow_schema::{DataType, Field, FieldRef, Schema};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion_common::{
@@ -685,6 +691,144 @@ impl Display for ArrayInsert {
     }
 }
 
+#[derive(Debug, Eq)]
+pub struct ArrayMinMax {
+    src_array_expr: Arc<dyn PhysicalExpr>,
+    legacy_negative_index: bool,
+    is_min: bool,
+}
+
+impl ArrayMinMax {
+    pub fn new(
+        src_array_expr: Arc<dyn PhysicalExpr>,
+        legacy_negative_index: bool,
+        is_min: bool,
+    ) -> Self {
+        Self {
+            src_array_expr,
+            legacy_negative_index,
+            is_min,
+        }
+    }
+    pub fn array_type(&self, data_type: &DataType) -> DataFusionResult<DataType> {
+        match data_type {
+            DataType::List(field) => Ok(field.data_type().clone()),
+            DataType::LargeList(field) => Ok(field.data_type().clone()),
+            data_type => Err(DataFusionError::Internal(format!(
+                "Unexpected src array type in ArrayMin / ArrayMax: {:?}",
+                data_type
+            ))),
+        }
+    }
+}
+
+impl Hash for ArrayMinMax {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.src_array_expr.hash(state);
+        self.legacy_negative_index.hash(state);
+        self.is_min.hash(state);
+    }
+}
+
+impl PartialEq for ArrayMinMax {
+    fn eq(&self, other: &Self) -> bool {
+        self.src_array_expr.eq(&other.src_array_expr)
+            && self.legacy_negative_index.eq(&other.legacy_negative_index)
+            && self.is_min.eq(&other.is_min)
+    }
+}
+
+impl Display for ArrayMinMax {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "ArrayMin / ArrayMax [array: {:?}, legacy_negative_index: {:?}]",
+            self.src_array_expr, self.legacy_negative_index
+        )
+    }
+}
+
+impl PhysicalExpr for ArrayMinMax {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn data_type(&self, input_schema: &Schema) -> DataFusionResult<DataType> {
+        self.array_type(&self.src_array_expr.data_type(input_schema)?)
+    }
+
+    fn nullable(&self, input_schema: &Schema) -> DataFusionResult<bool> {
+        self.src_array_expr.nullable(input_schema)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
+        let array = self
+            .src_array_expr
+            .evaluate(batch)?
+            .into_array(batch.num_rows())?;
+
+        let array_element_type = self.array_type(array.data_type())?;
+
+        match array.data_type() {
+            DataType::List(_) if array_element_type.is_primitive() => {
+                let list_array = as_list_array(&array)?;
+                array_min_max::<Int32Type, _>(&list_array, &array_element_type, self.is_min)
+            }
+            DataType::LargeList(_) if array_element_type.is_primitive() => {
+                let list_array = as_large_list_array(&array)?;
+                array_min_max::<Int32Type, _>(&list_array, &array_element_type, self.is_min)
+            }
+            _ => unimplemented!("ArrayMin / ArrayMax is not implemented yet"),
+        }
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        vec![&self.src_array_expr]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
+        Ok(Arc::new(ArrayMinMax::new(
+            Arc::clone(&children[0]),
+            self.legacy_negative_index,
+            self.is_min,
+        )))
+    }
+}
+
+fn array_min_max<T, O: OffsetSizeTrait>(
+    list_array: &GenericListArray<O>,
+    _element_type: &DataType,
+    is_min: bool,
+) -> DataFusionResult<ColumnarValue>
+where
+    T: ArrowNumericType,
+{
+    let mut builder = PrimitiveBuilder::<T>::with_capacity(list_array.len());
+    for array in list_array.iter() {
+        match array {
+            Some(a) => {
+                let primitive_array = a.as_primitive::<T>();
+                let value = if is_min {
+                    compute::min(primitive_array)
+                } else {
+                    compute::max(primitive_array)
+                };
+                match value {
+                    Some(v) => builder.append_value(v),
+                    None => builder.append_null(),
+                }
+            }
+            None => {
+                builder.append_null();
+            }
+        }
+    }
+    Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+}
+
 #[cfg(test)]
 mod test {
     use crate::list::{array_insert, list_extract, zero_based_index};
@@ -847,4 +991,14 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_array_min() -> Result<()> {
+        unimplemented!()
+    }
+
+    #[test]
+    fn test_array_max() -> Result<()> {
+        unimplemented!()
+    }
 }

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -2278,6 +2278,34 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
             expr.children(1),
             inputs,
             (builder, binaryExpr) => builder.setArrayAppend(binaryExpr))
+        case _ if expr.prettyName == "array_min" =>
+          createUnaryExpr(
+            expr.children.head,
+            inputs,
+            (builder, unaryExpr) => builder.setArrayMin(unaryExpr))
+        case _ if expr.prettyName == "array_max" =>
+          createUnaryExpr(
+            expr.children.head,
+            inputs,
+            (builder, unaryExpr) => builder.setArrayMax(unaryExpr))
+        case _ if expr.prettyName == "sort_array" =>
+          createBinaryExpr(
+            expr.children(0),
+            expr.children(1),
+            inputs,
+            (builder, binaryExpr) => builder.setSortArray(binaryExpr))
+        case _ if expr.prettyName == "array_zip" =>
+          createBinaryExpr(
+            expr.children(0),
+            expr.children(1),
+            inputs,
+            (builder, binaryExpr) => builder.setArrayZip(binaryExpr))
+        case _ if expr.prettyName == "array_union" =>
+          createBinaryExpr(
+            expr.children(0),
+            expr.children(1),
+            inputs,
+            (builder, binaryExpr) => builder.setArrayUnion(binaryExpr))
         case _ =>
           withInfo(expr, s"${expr.prettyName} is not supported", expr.children: _*)
           None

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -2529,4 +2529,27 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         spark.sql("SELECT array_contains((CASE WHEN _2 =_3 THEN array(_4) END), _4) FROM t1"));
     }
   }
+
+  test("array_union") {
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "test.parquet")
+      makeParquetFileAllTypes(path, dictionaryEnabled = false, n = 10000)
+      spark.read.parquet(path.toString).createOrReplaceTempView("t1");
+      checkSparkAnswerAndOperator(
+        spark.sql("SELECT array_union(array(_2, _3, _4), array(_2, _3, _4)) FROM t1"))
+    }
+  }
+
+  test("sort_array") {
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "test.parquet")
+      makeParquetFileAllTypes(path, dictionaryEnabled = false, n = 10000)
+      spark.read.parquet(path.toString).createOrReplaceTempView("t1");
+      val q = spark.sql("SELECT sort_array(array(_2, _3, _4)) FROM t1");
+      q.explain(true)
+      q.show(false)
+      // FIXME: Why collect fails ?  could not cast value to arrow_array::array::byte_array::GenericByteArray<arrow_array::types::GenericStringType<i32>>
+      // q.collect();
+    }
+  }
 }