Skip to content

Commit cf1c43a

Browse files
crm26claude
crm26
and
claude
authored
feat: add array_avg scalar function (#23168)
## Summary Adds `array_avg`, the last function from the split-PR pipeline tracked in #21536. Computes the arithmetic mean (sum divided by count) of the elements of a numeric array, returned as `Float64`. Templated from the merged `array_sum` (#22542) with the same SQL aggregate NULL conventions; sibling of `array_sum`, `array_product`, `array_subtract`, `array_add`, `array_scale`, and `array_normalize`. ## Semantics **NULL semantics — SQL aggregate convention (deliberate divergence from binary-op siblings):** - NULL row → NULL row out - NULL elements are **skipped** from BOTH the sum and the count, matching PostgreSQL `AVG`, DuckDB `list_avg`, Spark `aggregate`. So `array_avg([1, NULL, 3]) = (1+3) / 2 = 2`, not `(1+3) / 3`. - All-NULL row → NULL out (matches `AVG(...)` over an all-NULL column) - **Empty array → NULL** (matches sibling `array_sum` #22542 and `array_product` #22703, PostgreSQL, DuckDB `list_avg`, SQL Standard AVG-of-empty-set) **Type coercion:** - Inner numeric types (`Float32`, `Int*`, `UInt*`) coerced to `Float64`. Integer-arg literals are coerced too. - Return type is always `Float64` (since avg of integers can be non-integer). **List shapes supported:** `List`, `LargeList`, `FixedSizeList`. **Alias:** `list_avg` (matches the `list_sum` / `list_product` pattern). ## Test coverage SLT (`array_avg.slt`) covers: - Happy paths (basic, single element, negative values, cancelling positive/negative, non-integer mean `[1,2] → 1.5`) - All NULL shapes: bare `NULL` row, NULL elements skipped, single non-NULL among NULLs, all-NULL array - Empty array → NULL - All 3 list shapes (`List`, `LargeList`, `FixedSizeList`) - Float32, Int64 inner types, integer literals, integer mean that is non-integer - Multi-row mix - Error paths (non-list input, zero args, two args) - Return type assertion (`Float64`) - `list_avg` alias ## Closes The last open slot in #21536 — completes the 8-function split-PR pipeline from the originally-too-big PRs #21371 / #21376. Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 77545a4 commit cf1c43a

4 files changed

Lines changed: 375 additions & 0 deletions

File tree

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

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
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_avg 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+
ArrayAvg,
39+
array_avg,
40+
array,
41+
"returns the arithmetic mean of elements in a numeric array.",
42+
array_avg_udf
43+
);
44+
45+
#[user_doc(
46+
doc_section(label = "Array Functions"),
47+
description = "Returns the arithmetic mean (sum divided by count) of the elements of the input array. NULL elements are skipped (per SQL aggregate convention) and excluded from the count. Returns NULL if the input row is NULL, every element is NULL, or the array is empty.",
48+
syntax_example = "array_avg(array)",
49+
sql_example = r#"```sql
50+
> select array_avg([1.0, 2.0, 3.0]);
51+
+----------------------------+
52+
| array_avg(List([1.0,2.0,3.0])) |
53+
+----------------------------+
54+
| 2.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 ArrayAvg {
64+
signature: Signature,
65+
aliases: Vec<String>,
66+
}
67+
68+
impl Default for ArrayAvg {
69+
fn default() -> Self {
70+
Self::new()
71+
}
72+
}
73+
74+
impl ArrayAvg {
75+
pub fn new() -> Self {
76+
Self {
77+
signature: Signature::user_defined(Volatility::Immutable),
78+
aliases: vec!["list_avg".to_string()],
79+
}
80+
}
81+
}
82+
83+
impl ScalarUDFImpl for ArrayAvg {
84+
fn name(&self) -> &str {
85+
"array_avg"
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_avg_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_avg_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
127+
let [array] = take_function_args("array_avg", args)?;
128+
match array.data_type() {
129+
List(_) => general_array_avg::<i32>(array),
130+
LargeList(_) => general_array_avg::<i64>(array),
131+
arg_type => {
132+
internal_err!("array_avg received unexpected type after coercion: {arg_type}")
133+
}
134+
}
135+
}
136+
137+
fn general_array_avg<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+
// AVG, DuckDB list_avg, Spark aggregate). Empty arrays and all-NULL
155+
// arrays both yield NULL — same behavior as SQL AVG over an empty
156+
// set or all-NULL column.
157+
let mut sum = 0.0_f64;
158+
let mut count: u64 = 0;
159+
for i in start..end {
160+
if values.is_valid(i) {
161+
sum += values.value(i);
162+
count += 1;
163+
}
164+
}
165+
166+
if count > 0 {
167+
builder.append_value(sum / count as f64);
168+
} else {
169+
builder.append_null();
170+
}
171+
}
172+
173+
Ok(Arc::new(builder.finish()))
174+
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ pub mod macros_lambda;
4242

4343
pub mod array_add;
4444
pub mod array_any_match;
45+
pub mod array_avg;
4546
pub mod array_compact;
4647
pub mod array_filter;
4748
pub mod array_has;
@@ -95,6 +96,7 @@ use std::sync::Arc;
9596
pub mod expr_fn {
9697
pub use super::array_add::array_add;
9798
pub use super::array_any_match::array_any_match;
99+
pub use super::array_avg::array_avg;
98100
pub use super::array_compact::array_compact;
99101
pub use super::array_filter::array_filter;
100102
pub use super::array_has::array_has;
@@ -181,6 +183,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
181183
length::array_length_udf(),
182184
array_normalize::array_normalize_udf(),
183185
array_add::array_add_udf(),
186+
array_avg::array_avg_udf(),
184187
array_product::array_product_udf(),
185188
array_scale::array_scale_udf(),
186189
array_subtract::array_subtract_udf(),

‎datafusion/sqllogictest/test_files/array_avg.slt‎

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

0 commit comments

Comments
 (0)