datafusion_functions_nested/
array_avg.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_avg 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    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)]
63pub struct ArrayAvg {
64    signature: Signature,
65    aliases: Vec<String>,
66}
67
68impl Default for ArrayAvg {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74impl 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
83impl 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
126fn 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
137fn 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/array_avg.rs

datafusion_functions_nested/
array_avg.rs
