ArraySum

datafusion/expr/src/built_in_function.rs

@@ -193,6 +193,10 @@ pub enum BuiltinScalarFunction {
     Flatten,
     /// Range
     Range,
+    /// Aggregate
+    ArrayAggregate,
+    /// Sum
+    ArraySum,
 
     // struct functions
     /// struct
@@ -400,6 +404,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::Tan => Volatility::Immutable,
             BuiltinScalarFunction::Tanh => Volatility::Immutable,
             BuiltinScalarFunction::Trunc => Volatility::Immutable,
+            BuiltinScalarFunction::ArrayAggregate => Volatility::Immutable,
             BuiltinScalarFunction::ArrayAppend => Volatility::Immutable,
             BuiltinScalarFunction::ArraySort => Volatility::Immutable,
             BuiltinScalarFunction::ArrayConcat => Volatility::Immutable,
@@ -427,6 +432,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayReplaceAll => Volatility::Immutable,
             BuiltinScalarFunction::Flatten => Volatility::Immutable,
             BuiltinScalarFunction::ArraySlice => Volatility::Immutable,
+            BuiltinScalarFunction::ArraySum => Volatility::Immutable,
             BuiltinScalarFunction::ArrayToString => Volatility::Immutable,
             BuiltinScalarFunction::ArrayIntersect => Volatility::Immutable,
             BuiltinScalarFunction::ArrayUnion => Volatility::Immutable,
@@ -549,6 +555,10 @@ impl BuiltinScalarFunction {
                 let data_type = get_base_type(&input_expr_types[0])?;
                 Ok(data_type)
             }
+
+            BuiltinScalarFunction::ArrayAggregate => {
+                internal_err!("ArrayAggregate should be rewritten to other function")
+            }
             BuiltinScalarFunction::ArrayAppend => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArraySort => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArrayConcat => {
@@ -617,6 +627,15 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrayReplaceN => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArrayReplaceAll => Ok(input_expr_types[0].clone()),
             BuiltinScalarFunction::ArraySlice => Ok(input_expr_types[0].clone()),
+            BuiltinScalarFunction::ArraySum => {
+                if let DataType::List(field) = &input_expr_types[0] {
+                    Ok(field.data_type().clone())
+                } else {
+                    plan_err!(
+                        "The {self} function can only accept list as the first argument"
+                    )
+                }
+            }
             BuiltinScalarFunction::ArrayToString => Ok(Utf8),
             BuiltinScalarFunction::ArrayUnion | BuiltinScalarFunction::ArrayIntersect => {
                 match (input_expr_types[0].clone(), input_expr_types[1].clone()) {
@@ -927,6 +946,9 @@ impl BuiltinScalarFunction {
                 // 0 or more arguments of arbitrary type
                 Signature::one_of(vec![VariadicEqual, Any(0)], self.volatility())
             }
+            BuiltinScalarFunction::ArrayAggregate => {
+                unreachable!("ArrayAggregate should be rewritten to other function")
+            }
             BuiltinScalarFunction::ArrayPopFront => Signature::any(1, self.volatility()),
             BuiltinScalarFunction::ArrayPopBack => Signature::any(1, self.volatility()),
             BuiltinScalarFunction::ArrayConcat => {
@@ -960,6 +982,7 @@ impl BuiltinScalarFunction {
                 Signature::any(3, self.volatility())
             }
             BuiltinScalarFunction::ArraySlice => Signature::any(3, self.volatility()),
+            BuiltinScalarFunction::ArraySum => Signature::any(1, self.volatility()),
             BuiltinScalarFunction::ArrayToString => {
                 Signature::variadic_any(self.volatility())
             }
@@ -1564,6 +1587,12 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::ArrowTypeof => &["arrow_typeof"],
 
             // array functions
+            BuiltinScalarFunction::ArrayAggregate => &[
+                "array_aggregate",
+                "list_aggregate",
+                "array_aggr",
+                "list_aggr",
+            ],
             BuiltinScalarFunction::ArrayAppend => &[
                 "array_append",
                 "list_append",
@@ -1625,12 +1654,14 @@ impl BuiltinScalarFunction {
                 &["array_replace_all", "list_replace_all"]
             }
             BuiltinScalarFunction::ArraySlice => &["array_slice", "list_slice"],
+            BuiltinScalarFunction::ArraySum => &["array_sum", "list_sum"],
             BuiltinScalarFunction::ArrayToString => &[
                 "array_to_string",
                 "list_to_string",
                 "array_join",
                 "list_join",
             ],
+
             BuiltinScalarFunction::ArrayUnion => &["array_union", "list_union"],
             BuiltinScalarFunction::Cardinality => &["cardinality"],
             BuiltinScalarFunction::MakeArray => &["make_array", "make_list"],

datafusion/expr/src/expr_fn.rs

@@ -576,6 +576,13 @@ scalar_expr!(Uuid, uuid, , "returns uuid v4 as a string value");
 scalar_expr!(Log, log, base x, "logarithm of a `x` for a particular `base`");
 
 // array functions
+scalar_expr!(
+    ArrayAggregate,
+    array_aggregate,
+    array name,
+    "allows the execution of arbitrary existing aggregate functions `name` on the elements of a list"
+);
+
 scalar_expr!(
     ArrayAppend,
     array_append,
@@ -732,6 +739,7 @@ scalar_expr!(
     array offset length,
     "returns a slice of the array."
 );
+scalar_expr!(ArraySum, array_sum, array, "return sum of the array");
 scalar_expr!(
     ArrayToString,
     array_to_string,

datafusion/physical-expr/src/array_aggregate.rs

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Array Aggregate expressions
+
+use std::sync::Arc;
+
+use arrow::compute;
+use arrow_array::{
+    cast::AsArray,
+    types::{Decimal128Type, Decimal256Type, Float64Type, Int64Type, UInt64Type},
+    Array, ArrayRef, ArrowNumericType, GenericListArray, OffsetSizeTrait, PrimitiveArray,
+};
+use arrow_buffer::NullBuffer;
+use arrow_schema::DataType;
+use datafusion_common::DataFusionError;
+use datafusion_common::Result;
+use datafusion_common::{cast::as_list_array, not_impl_err};
+
+/// ArraySum aggregate expression
+pub fn array_sum(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let list_array = as_list_array(&args[0])?;
+    match list_array.value_type() {
+        DataType::Int64 => {
+            Ok(Arc::new(array_sum_internal::<Int64Type, i32>(list_array)?))
+        }
+        DataType::UInt64 => {
+            Ok(Arc::new(array_sum_internal::<UInt64Type, i32>(list_array)?))
+        }
+        DataType::Float64 => Ok(Arc::new(array_sum_internal::<Float64Type, i32>(
+            list_array,
+        )?)),
+        DataType::Decimal128(_, _) => Ok(Arc::new(array_sum_internal::<
+            Decimal128Type,
+            i32,
+        >(list_array)?)),
+        DataType::Decimal256(_, _) => Ok(Arc::new(array_sum_internal::<
+            Decimal256Type,
+            i32,
+        >(list_array)?)),
+        _ => not_impl_err!(
+            "array_sum for type {:?} not implemented",
+            list_array.value_type()
+        ),
+    }
+}
+
+fn array_sum_internal<T: ArrowNumericType, O: OffsetSizeTrait>(
+    array: &GenericListArray<O>,
+) -> Result<PrimitiveArray<T>> {
+    let mut data = Vec::with_capacity(array.len());
+    for arr in array.iter() {
+        if let Some(arr) = arr {
+            let i64arr = arr.as_primitive::<T>();
+            let sum = compute::sum_checked(i64arr)?;
+            data.push(sum);
+        } else {
+            data.push(None);
+        }
+    }
+
+    let mut values = Vec::with_capacity(array.len());
+    let mut nulls = Vec::with_capacity(array.len());
+    for v in data.into_iter() {
+        if let Some(v) = v {
+            values.push(v);
+            nulls.push(true);
+        } else {
+            values.push(T::default_value());
+            nulls.push(false);
+        }
+    }
+
+    Ok(PrimitiveArray::<T>::new(
+        values.into(),
+        Some(NullBuffer::from(nulls)),
+    ))
+}

datafusion/physical-expr/src/functions.rs

@@ -33,7 +33,7 @@
 use crate::execution_props::ExecutionProps;
 use crate::sort_properties::SortProperties;
 use crate::{
-    array_expressions, conditional_expressions, datetime_expressions,
+    array_aggregate, array_expressions, conditional_expressions, datetime_expressions,
     expressions::nullif_func, math_expressions, string_expressions, struct_expressions,
     PhysicalExpr, ScalarFunctionExpr,
 };
@@ -326,6 +326,9 @@ pub fn create_physical_fun(
         }
 
         // array functions
+        BuiltinScalarFunction::ArrayAggregate => {
+            unreachable!("ArrayAggregate should be rewritten to other function")
+        }
         BuiltinScalarFunction::ArrayAppend => {
             Arc::new(|args| make_scalar_function(array_expressions::array_append)(args))
         }
@@ -407,6 +410,9 @@ pub fn create_physical_fun(
         BuiltinScalarFunction::ArraySlice => {
             Arc::new(|args| make_scalar_function(array_expressions::array_slice)(args))
         }
+        BuiltinScalarFunction::ArraySum => {
+            Arc::new(|args| make_scalar_function(array_aggregate::array_sum)(args))
+        }
         BuiltinScalarFunction::ArrayToString => Arc::new(|args| {
             make_scalar_function(array_expressions::array_to_string)(args)
         }),

datafusion/physical-expr/src/lib.rs

@@ -17,6 +17,7 @@
 
 pub mod aggregate;
 pub mod analysis;
+pub mod array_aggregate;
 pub mod array_expressions;
 pub mod conditional_expressions;
 #[cfg(feature = "crypto_expressions")]

datafusion/proto/proto/datafusion.proto

@@ -667,6 +667,9 @@ enum ScalarFunction {
   FindInSet = 127;
   ArraySort = 128;
   ArrayDistinct = 129;
+
+  ArrayAggregate = 200;
+  ArraySum = 201;
 }
 
 message ScalarFunctionNode {

datafusion/proto/src/logical_plan/from_proto.rs

@@ -42,15 +42,15 @@ use datafusion_common::{
 };
 use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by};
 use datafusion_expr::{
-    abs, acos, acosh, array, array_append, array_concat, array_dims, array_distinct,
-    array_element, array_except, array_has, array_has_all, array_has_any,
+    abs, acos, acosh, array, array_aggregate, array_append, array_concat, array_dims,
+    array_distinct, array_element, array_except, array_has, array_has_all, array_has_any,
     array_intersect, array_length, array_ndims, array_position, array_positions,
     array_prepend, array_remove, array_remove_all, array_remove_n, array_repeat,
     array_replace, array_replace_all, array_replace_n, array_slice, array_sort,
-    array_to_string, arrow_typeof, ascii, asin, asinh, atan, atan2, atanh, bit_length,
-    btrim, cardinality, cbrt, ceil, character_length, chr, coalesce, concat_expr,
-    concat_ws_expr, cos, cosh, cot, current_date, current_time, date_bin, date_part,
-    date_trunc, decode, degrees, digest, encode, exp,
+    array_sum, array_to_string, arrow_typeof, ascii, asin, asinh, atan, atan2, atanh,
+    bit_length, btrim, cardinality, cbrt, ceil, character_length, chr, coalesce,
+    concat_expr, concat_ws_expr, cos, cosh, cot, current_date, current_time, date_bin,
+    date_part, date_trunc, decode, degrees, digest, encode, exp,
     expr::{self, InList, Sort, WindowFunction},
     factorial, find_in_set, flatten, floor, from_unixtime, gcd, gen_range, isnan, iszero,
     lcm, left, levenshtein, ln, log, log10, log2,
@@ -476,6 +476,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::ArraySort => Self::ArraySort,
             ScalarFunction::ArrayConcat => Self::ArrayConcat,
@@ -503,6 +504,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
             ScalarFunction::ArrayReplaceN => Self::ArrayReplaceN,
             ScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll,
             ScalarFunction::ArraySlice => Self::ArraySlice,
+            ScalarFunction::ArraySum => Self::ArraySum,
             ScalarFunction::ArrayToString => Self::ArrayToString,
             ScalarFunction::ArrayIntersect => Self::ArrayIntersect,
             ScalarFunction::ArrayUnion => Self::ArrayUnion,
@@ -1361,6 +1363,10 @@ pub fn parse_expr(
                         .map(|expr| parse_expr(expr, registry))
                         .collect::<Result<Vec<_>, _>>()?,
                 )),
+                ScalarFunction::ArrayAggregate => Ok(array_aggregate(
+                    parse_expr(&args[0], registry)?,
+                    parse_expr(&args[1], registry)?,
+                )),
                 ScalarFunction::ArrayAppend => Ok(array_append(
                     parse_expr(&args[0], registry)?,
                     parse_expr(&args[1], registry)?,
@@ -1453,6 +1459,9 @@ pub fn parse_expr(
                     parse_expr(&args[1], registry)?,
                     parse_expr(&args[2], registry)?,
                 )),
+                ScalarFunction::ArraySum => {
+                    Ok(array_sum(parse_expr(&args[0], registry)?))
+                }
                 ScalarFunction::ArrayToString => Ok(array_to_string(
                     parse_expr(&args[0], registry)?,
                     parse_expr(&args[1], registry)?,

datafusion/proto/src/logical_plan/to_proto.rs

@@ -1503,6 +1503,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::ArraySort => Self::ArraySort,
             BuiltinScalarFunction::ArrayConcat => Self::ArrayConcat,
@@ -1530,6 +1531,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
             BuiltinScalarFunction::ArrayReplaceN => Self::ArrayReplaceN,
             BuiltinScalarFunction::ArrayReplaceAll => Self::ArrayReplaceAll,
             BuiltinScalarFunction::ArraySlice => Self::ArraySlice,
+            BuiltinScalarFunction::ArraySum => Self::ArraySum,
             BuiltinScalarFunction::ArrayToString => Self::ArrayToString,
             BuiltinScalarFunction::ArrayIntersect => Self::ArrayIntersect,
             BuiltinScalarFunction::ArrayUnion => Self::ArrayUnion,