add aggregate

Signed-off-by: jayzhan211 <[email protected]>

Author: jayzhan211

datafusion/core/tests/sqllogictests/test_files/array.slt

@@ -2332,6 +2332,12 @@ select array_concat(column1, [7]) from arrays_values_v2;
 
 # array aggregate function
 
+## array aggregate
+query I
+select array_aggregate([1, 3, 5, 7], 'sum');
+----
+16
+
 ## array sum
 query IRI
 select array_sum([1, 3, 5, 7]),
@@ -2340,6 +2346,10 @@ select array_sum([1, 3, 5, 7]),
 ----
 16 6.6 23
 
+# TODO: Support nulls in array.
+# query error DataFusion error: This feature is not implemented: Arrays with different types are not supported: \{Int64, Null\}
+# select array_sum([1, null, 3, null]);
+
 query ??
 select column1, column6 from arrays_values_without_nulls;
 ----

datafusion/expr/src/built_in_function.rs

@@ -117,6 +117,8 @@ pub enum BuiltinScalarFunction {
     Cot,
 
     // array functions
+    /// array_aggregate
+    ArrayAggregate,
     /// array_append
     ArrayAppend,
     /// array_concat
@@ -350,6 +352,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::Tanh => Volatility::Immutable,
             BuiltinScalarFunction::Trunc => Volatility::Immutable,
             BuiltinScalarFunction::ArrayAppend => Volatility::Immutable,
+            BuiltinScalarFunction::ArrayAggregate => Volatility::Immutable,
             BuiltinScalarFunction::ArrayConcat => Volatility::Immutable,
             BuiltinScalarFunction::ArrayHasAll => Volatility::Immutable,
             BuiltinScalarFunction::ArrayHasAny => Volatility::Immutable,
@@ -503,6 +506,10 @@ impl BuiltinScalarFunction {
         // Some built-in functions' return type depends on the incoming type.
         match self {
             BuiltinScalarFunction::ArrayAppend => Ok(input_expr_types[0].clone()),
+            BuiltinScalarFunction::ArrayAggregate => {
+                // TODO: Fix this
+                Ok(Int64)
+            }
             BuiltinScalarFunction::ArrayConcat => {
                 let mut expr_type = Null;
                 let mut max_dims = 0;
@@ -839,15 +846,16 @@ impl BuiltinScalarFunction {
 
         // for now, the list is small, as we do not have many built-in functions.
         match self {
-            BuiltinScalarFunction::ArrayAppend => Signature::any(2, self.volatility()),
             BuiltinScalarFunction::ArrayConcat => {
                 Signature::variadic_any(self.volatility())
             }
-            BuiltinScalarFunction::ArrayHasAll
+            BuiltinScalarFunction::ArrayAggregate
+            | BuiltinScalarFunction::ArrayAppend
+            | BuiltinScalarFunction::ArrayElement
+            | BuiltinScalarFunction::ArrayHasAll
             | BuiltinScalarFunction::ArrayHasAny
             | BuiltinScalarFunction::ArrayHas => Signature::any(2, self.volatility()),
             BuiltinScalarFunction::ArrayDims => Signature::any(1, self.volatility()),
-            BuiltinScalarFunction::ArrayElement => Signature::any(2, self.volatility()),
             BuiltinScalarFunction::ArrayLength => {
                 Signature::variadic_any(self.volatility())
             }
@@ -1316,6 +1324,12 @@ fn aliases(func: &BuiltinScalarFunction) -> &'static [&'static str] {
         BuiltinScalarFunction::ArrowTypeof => &["arrow_typeof"],
 
         // array functions
+        BuiltinScalarFunction::ArrayAggregate => &[
+            "array_aggregate",
+            "list_aggregate",
+            "array_aggr",
+            "list_aggr",
+        ],
         BuiltinScalarFunction::ArrayAppend => &[
             "array_append",
             "list_append",

datafusion/expr/src/expr_fn.rs

@@ -545,6 +545,12 @@ scalar_expr!(
     array element,
     "appends an element to the end of an array."
 );
+scalar_expr!(
+    ArrayAggregate,
+    array_aggregate,
+    array name,
+    "allows the execution of arbitrary existing aggregate functions on the elements of a list"
+);
 nary_scalar_expr!(ArrayConcat, array_concat, "concatenates arrays.");
 scalar_expr!(
     ArrayHas,

datafusion/physical-expr/src/array_expressions.rs

@@ -2117,7 +2117,20 @@ macro_rules! array_sum_internal(
     }}
 );
 
-/// Array_sum SQL function
+/// array_aggregate SQL function
+pub fn array_aggregate(args: &[ArrayRef]) -> Result<ArrayRef> {
+    assert_eq!(args.len(), 2);
+    let func_name = args[1].as_string::<i32>().value(0);
+    let args = &args[0..1];
+    match func_name {
+        "sum" => array_sum(args),
+        _ => Err(DataFusionError::NotImplemented(format!(
+            "array_aggregate does not support function '{func_name}'"
+        ))),
+    }
+}
+
+/// array_sum SQL function
 pub fn array_sum(args: &[ArrayRef]) -> Result<ArrayRef> {
     assert_eq!(args.len(), 1);
 

datafusion/physical-expr/src/functions.rs

@@ -416,6 +416,9 @@ pub fn create_physical_fun(
         BuiltinScalarFunction::ArrayAppend => {
             Arc::new(|args| make_scalar_function(array_expressions::array_append)(args))
         }
+        BuiltinScalarFunction::ArrayAggregate => Arc::new(|args| {
+            make_scalar_function(array_expressions::array_aggregate)(args)
+        }),
         BuiltinScalarFunction::ArrayConcat => {
             Arc::new(|args| make_scalar_function(array_expressions::array_concat)(args))
         }

datafusion/proto/proto/datafusion.proto

@@ -595,6 +595,7 @@ enum ScalarFunction {
   ArrayReplaceAll = 110;
   Nanvl = 111;
   ArraySum = 113;
+  ArrayAggregate = 114;
 }
 
 message ScalarFunctionNode {

datafusion/proto/src/generated/pbjson.rs

@@ -450,6 +450,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
             ScalarFunction::Ltrim => Self::Ltrim,
             ScalarFunction::Rtrim => Self::Rtrim,
             ScalarFunction::ToTimestamp => Self::ToTimestamp,
+            ScalarFunction::ArrayAggregate => Self::ArrayAggregate,
             ScalarFunction::ArrayAppend => Self::ArrayAppend,
             ScalarFunction::ArrayConcat => Self::ArrayConcat,
             ScalarFunction::ArrayHasAll => Self::ArrayHasAll,

datafusion/proto/src/generated/prost.rs

@@ -1449,6 +1449,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
             BuiltinScalarFunction::Ltrim => Self::Ltrim,
             BuiltinScalarFunction::Rtrim => Self::Rtrim,
             BuiltinScalarFunction::ToTimestamp => Self::ToTimestamp,
+            BuiltinScalarFunction::ArrayAggregate => Self::ArrayAggregate,
             BuiltinScalarFunction::ArrayAppend => Self::ArrayAppend,
             BuiltinScalarFunction::ArrayConcat => Self::ArrayConcat,
             BuiltinScalarFunction::ArrayHasAll => Self::ArrayHasAll,