datafusion_functions_nested/
array_sum.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`ScalarUDFImpl`] definitions for array_sum function.
19
20use crate::utils::make_scalar_function;
21use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
22use arrow::datatypes::{
23    DataType,
24    DataType::{FixedSizeList, LargeList, List, Null},
25    Field,
26};
27use datafusion_common::cast::{as_float64_array, as_generic_list_array};
28use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
29use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args};
30use datafusion_expr::{
31    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    Volatility,
33};
34use datafusion_macros::user_doc;
35use std::sync::Arc;
36
37make_udf_expr_and_func!(
38    ArraySum,
39    array_sum,
40    array,
41    "returns the sum of elements in a numeric array.",
42    array_sum_udf
43);
44
45#[user_doc(
46    doc_section(label = "Array Functions"),
47    description = "Returns the sum of the elements of the input array, computed as `array[0] + array[1] + ...`. NULL elements are skipped (per SQL aggregate convention). Returns NULL if the input row is NULL, every element is NULL, or the array is empty.",
48    syntax_example = "array_sum(array)",
49    sql_example = r#"```sql
50> select array_sum([1.0, 2.0, 3.0]);
51+----------------------------+
52| array_sum(List([1.0,2.0,3.0])) |
53+----------------------------+
54| 6.0                        |
55+----------------------------+
56```"#,
57    argument(
58        name = "array",
59        description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
60    )
61)]
62#[derive(Debug, PartialEq, Eq, Hash)]
63pub struct ArraySum {
64    signature: Signature,
65    aliases: Vec<String>,
66}
67
68impl Default for ArraySum {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74impl ArraySum {
75    pub fn new() -> Self {
76        Self {
77            signature: Signature::user_defined(Volatility::Immutable),
78            aliases: vec!["list_sum".to_string()],
79        }
80    }
81}
82
83impl ScalarUDFImpl for ArraySum {
84    fn name(&self) -> &str {
85        "array_sum"
86    }
87
88    fn signature(&self) -> &Signature {
89        &self.signature
90    }
91
92    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
93        Ok(DataType::Float64)
94    }
95
96    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
97        let [arg_type] = take_function_args(self.name(), arg_types)?;
98        let coercion = Some(&ListCoercion::FixedSizedListToList);
99
100        if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) {
101            return plan_err!("{} does not support type {arg_type}", self.name());
102        }
103
104        let coerced = if matches!(arg_type, Null) {
105            List(Arc::new(Field::new_list_field(DataType::Float64, true)))
106        } else {
107            coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion)
108        };
109
110        Ok(vec![coerced])
111    }
112
113    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
114        make_scalar_function(array_sum_inner)(&args.args)
115    }
116
117    fn aliases(&self) -> &[String] {
118        &self.aliases
119    }
120
121    fn documentation(&self) -> Option<&Documentation> {
122        self.doc()
123    }
124}
125
126fn array_sum_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
127    let [array] = take_function_args("array_sum", args)?;
128    match array.data_type() {
129        List(_) => general_array_sum::<i32>(array),
130        LargeList(_) => general_array_sum::<i64>(array),
131        arg_type => {
132            internal_err!("array_sum received unexpected type after coercion: {arg_type}")
133        }
134    }
135}
136
137fn general_array_sum<O: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
138    let list_array = as_generic_list_array::<O>(array)?;
139    let values = as_float64_array(list_array.values())?;
140    let offsets = list_array.value_offsets();
141
142    let mut builder = Float64Array::builder(list_array.len());
143
144    for row in 0..list_array.len() {
145        if list_array.is_null(row) {
146            builder.append_null();
147            continue;
148        }
149
150        let start = offsets[row].as_usize();
151        let end = offsets[row + 1].as_usize();
152
153        // Skip NULL elements per SQL aggregate convention (matches PostgreSQL
154        // array_sum, DuckDB list_sum, Spark aggregate). Empty arrays and
155        // all-NULL arrays both yield NULL — same behavior as SQL SUM over
156        // an empty set or all-NULL column.
157        let mut sum = 0.0_f64;
158        let mut any_valid = false;
159        for i in start..end {
160            if values.is_valid(i) {
161                sum += values.value(i);
162                any_valid = true;
163            }
164        }
165
166        if any_valid {
167            builder.append_value(sum);
168        } else {
169            builder.append_null();
170        }
171    }
172
173    Ok(Arc::new(builder.finish()))
174}
datafusion_functions_nested/array_sum.rs

datafusion_functions_nested/
array_sum.rs
