Skip to content

Commit d4049bf

Browse files
crm26crm26
committed
feat: add array_sum scalar function
Adds `array_sum(array)` returning the sum of elements in a numeric array. Aliased as `list_sum`. Part of the per-function split sequence on tracking issue apache#21536, following the pattern of the already-merged PRs in this series (cosine_distance apache#21542, inner_product apache#21861, array_normalize apache#22013, array_scale apache#22466). Semantics: - NULL row in array -> NULL row out - NULL elements are skipped (SQL aggregate convention; matches PostgreSQL array_sum, DuckDB list_sum, Spark aggregate). A row whose every element is NULL yields NULL. - Empty array -> 0.0 (additive identity, matches SQL SUM over no rows conceptually, and DuckDB list_sum([]) = 0) Input is List/LargeList/FixedSizeList of any numeric type; elements are coerced to Float64. Output is Float64.
1 parent 8036b94 commit d4049bf

3 files changed

Lines changed: 339 additions & 0 deletions

File tree

‎datafusion/functions-nested/src/array_sum.rs‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
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+
20+
use crate::utils::make_scalar_function;
21+
use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
22+
use arrow::datatypes::{
23+
DataType,
24+
DataType::{FixedSizeList, LargeList, List, Null},
25+
Field,
26+
};
27+
use datafusion_common::cast::{as_float64_array, as_generic_list_array};
28+
use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only};
29+
use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args};
30+
use datafusion_expr::{
31+
ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32+
Volatility,
33+
};
34+
use datafusion_macros::user_doc;
35+
use std::sync::Arc;
36+
37+
make_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 or every element is NULL. Returns 0.0 for an empty array.",
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)]
63+
pub struct ArraySum {
64+
signature: Signature,
65+
aliases: Vec<String>,
66+
}
67+
68+
impl Default for ArraySum {
69+
fn default() -> Self {
70+
Self::new()
71+
}
72+
}
73+
74+
impl 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+
83+
impl 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+
126+
fn 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+
137+
fn 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+
let len = end - start;
153+
154+
// Empty array: sum is the additive identity. Matches SQL SUM(<empty>) = 0
155+
// and DuckDB's list_sum(([]) = 0 conventions.
156+
if len == 0 {
157+
builder.append_value(0.0);
158+
continue;
159+
}
160+
161+
// `slice` resets the logical offset to 0, so `i` below is 0-based within the slice.
162+
let slice = values.slice(start, len);
163+
164+
// Skip NULL elements per SQL aggregate convention (matches PostgreSQL
165+
// array_sum, DuckDB list_sum, Spark aggregate). A row with every
166+
// element NULL yields NULL — same behavior as SQL SUM over all-NULL.
167+
let mut sum = 0.0_f64;
168+
let mut any_valid = false;
169+
for i in 0..len {
170+
if !slice.is_null(i) {
171+
sum += slice.value(i);
172+
any_valid = true;
173+
}
174+
}
175+
176+
if any_valid {
177+
builder.append_value(sum);
178+
} else {
179+
builder.append_null();
180+
}
181+
}
182+
183+
Ok(Arc::new(builder.finish()))
184+
}

‎datafusion/functions-nested/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ pub mod array_normalize;
4949
pub mod array_product;
5050
pub mod array_scale;
5151
pub mod array_subtract;
52+
pub mod array_sum;
5253
pub mod array_transform;
5354
pub mod arrays_zip;
5455
pub mod cardinality;
@@ -103,6 +104,7 @@ pub mod expr_fn {
103104
pub use super::array_product::array_product;
104105
pub use super::array_scale::array_scale;
105106
pub use super::array_subtract::array_subtract;
107+
pub use super::array_sum::array_sum;
106108
pub use super::array_transform::array_transform;
107109
pub use super::arrays_zip::arrays_zip;
108110
pub use super::cardinality::cardinality;
@@ -182,6 +184,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
182184
array_product::array_product_udf(),
183185
array_scale::array_scale_udf(),
184186
array_subtract::array_subtract_udf(),
187+
array_sum::array_sum_udf(),
185188
cosine_distance::cosine_distance_udf(),
186189
inner_product::inner_product_udf(),
187190
distance::array_distance_udf(),

‎datafusion/sqllogictest/test_files/array_sum.slt‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
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+
## array_sum
19+
20+
# Basic case
21+
query R
22+
select array_sum([1.0, 2.0, 3.0]);
23+
----
24+
6
25+
26+
# Single element
27+
query R
28+
select array_sum([5.0]);
29+
----
30+
5
31+
32+
# Negative values
33+
query R
34+
select array_sum([-1.0, -2.0, -3.0]);
35+
----
36+
-6
37+
38+
# Positive and negative cancel
39+
query R
40+
select array_sum([1.0, -1.0, 2.0, -2.0]);
41+
----
42+
0
43+
44+
# Empty array returns 0 (additive identity, per SQL SUM convention)
45+
query R
46+
select array_sum(arrow_cast(make_array(), 'List(Float64)'));
47+
----
48+
0
49+
50+
# Bare NULL input returns NULL row
51+
query R
52+
select array_sum(NULL);
53+
----
54+
NULL
55+
56+
# NULL elements are skipped (SQL aggregate convention)
57+
query R
58+
select array_sum([1.0, NULL, 3.0]);
59+
----
60+
4
61+
62+
# Single NULL among numeric: skip the NULL
63+
query R
64+
select array_sum([NULL, 10.0]);
65+
----
66+
10
67+
68+
# All-NULL array returns NULL row (matches SQL SUM over all-NULL)
69+
query R
70+
select array_sum(arrow_cast([NULL, NULL], 'List(Float64)'));
71+
----
72+
NULL
73+
74+
# LargeList support
75+
query R
76+
select array_sum(arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)'));
77+
----
78+
6
79+
80+
# FixedSizeList input (coerced to List)
81+
query R
82+
select array_sum(arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)'));
83+
----
84+
6
85+
86+
# Float32 inner type (coerced to Float64)
87+
query R
88+
select array_sum(arrow_cast([1.0, 2.0, 3.0], 'List(Float32)'));
89+
----
90+
6
91+
92+
# Int64 inner type (coerced to Float64)
93+
query R
94+
select array_sum(arrow_cast([1, 2, 3], 'List(Int64)'));
95+
----
96+
6
97+
98+
# Integer literals (coerced to Float64)
99+
query R
100+
select array_sum([1, 2, 3]);
101+
----
102+
6
103+
104+
# Unsupported non-list input (plan error)
105+
query error array_sum does not support type
106+
select array_sum(1);
107+
108+
# Multi-row query with mix of normal, single-element, NULL elements, empty, NULL row
109+
query R
110+
select array_sum(column1) from (values
111+
(make_array(1.0, 2.0, 3.0)),
112+
(make_array(0.0)),
113+
(make_array(1.0, NULL, 4.0)),
114+
(arrow_cast(make_array(), 'List(Float64)')),
115+
(NULL)
116+
) as t(column1);
117+
----
118+
6
119+
0
120+
5
121+
0
122+
NULL
123+
124+
# Wrong arity (zero args)
125+
query error array_sum function requires 1 argument, got 0
126+
select array_sum();
127+
128+
# Wrong arity (two args)
129+
query error array_sum function requires 1 argument, got 2
130+
select array_sum([1.0], [2.0]);
131+
132+
# Return type is Float64
133+
query RT
134+
select array_sum([1.0, 2.0, 3.0]), arrow_typeof(array_sum([1.0, 2.0, 3.0]));
135+
----
136+
6 Float64
137+
138+
# list_sum alias produces the same result
139+
query R
140+
select list_sum([1.0, 2.0, 3.0]);
141+
----
142+
6
143+
144+
# list_sum alias with NULL row propagates correctly
145+
query R
146+
select list_sum(column1) from (values
147+
(make_array(1.0, 2.0)),
148+
(NULL)
149+
) as t(column1);
150+
----
151+
3
152+
NULL

0 commit comments

Comments
 (0)