From ad452218738124e3998b9393f990db17aa200ba4 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 10 Sep 2026 14:10:06 +0300 Subject: [PATCH] feat: add vectorized array_reduce --- datafusion/functions-nested/Cargo.toml | 4 + .../functions-nested/benches/array_reduce.rs | 122 +++++++ .../functions-nested/src/array_reduce.rs | 340 ++++++++++++++++++ datafusion/functions-nested/src/lib.rs | 3 + .../test_files/array/array_reduce.slt | 98 +++++ 5 files changed, 567 insertions(+) create mode 100644 datafusion/functions-nested/benches/array_reduce.rs create mode 100644 datafusion/functions-nested/src/array_reduce.rs create mode 100644 datafusion/sqllogictest/test_files/array/array_reduce.slt diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index ed5a89b8e3e72..ac793f42e490b 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -78,6 +78,10 @@ name = "array_concat" harness = false name = "array_min_max" +[[bench]] +harness = false +name = "array_reduce" + [[bench]] harness = false name = "arrays_zip" diff --git a/datafusion/functions-nested/benches/array_reduce.rs b/datafusion/functions-nested/benches/array_reduce.rs new file mode 100644 index 0000000000000..ed5e73af34df6 --- /dev/null +++ b/datafusion/functions-nested/benches/array_reduce.rs @@ -0,0 +1,122 @@ +// 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. + +use std::{collections::HashMap, sync::Arc}; + +use arrow::{ + array::{Array, ArrayRef, Int64Array, ListArray, RecordBatch}, + buffer::OffsetBuffer, + datatypes::{DataType, Field}, +}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::DFSchema; +use datafusion_expr::{ + Expr, col, + execution_props::ExecutionProps, + expr::{HigherOrderFunction, LambdaVariable}, + lambda, lit, + physical_planning_context::PhysicalPlanningContext, +}; +use datafusion_functions_nested::array_reduce::array_reduce_higher_order_function; +use datafusion_physical_expr::create_physical_expr; + +const NUM_ROWS: usize = 8192; +const LIST_SIZE: usize = 16; + +fn list_array(lengths: impl IntoIterator) -> ListArray { + let lengths = lengths.into_iter().collect::>(); + let values_len = lengths.iter().sum(); + let values = [1_i64, 2, 3, 4] + .into_iter() + .cycle() + .take(values_len) + .collect::>(); + ListArray::new( + Arc::new(Field::new_list_field(DataType::Int64, false)), + OffsetBuffer::from_lengths(lengths), + Arc::new(Int64Array::from(values)), + None, + ) +} + +fn reduce_expression( + list: &ListArray, +) -> (Arc, RecordBatch) { + let schema = DFSchema::from_unqualified_fields( + vec![Field::new("list", list.data_type().clone(), false)].into(), + HashMap::new(), + ) + .unwrap(); + let accumulator = Expr::LambdaVariable(LambdaVariable::new( + "acc".to_string(), + Some(Arc::new(Field::new("acc", DataType::Int64, false))), + )); + let value = Expr::LambdaVariable(LambdaVariable::new( + "value".to_string(), + Some(Arc::new(Field::new("value", DataType::Int64, false))), + )); + let expression = Expr::HigherOrderFunction(HigherOrderFunction::new( + array_reduce_higher_order_function(), + vec![ + col("list"), + lit(0_i64), + lambda(["acc", "value"], accumulator + value), + ], + )); + let physical = create_physical_expr( + &expression, + &schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + Arc::clone(schema.inner()), + vec![Arc::new(list.clone()) as ArrayRef], + ) + .unwrap(); + (physical, batch) +} + +fn criterion_benchmark(c: &mut Criterion) { + let inputs = [ + ( + "uniform", + list_array(std::iter::repeat_n(LIST_SIZE, NUM_ROWS)), + ), + ( + "varying", + list_array((0..NUM_ROWS).map(|row| match row % 8 { + 0 => 0, + 1 | 2 => 1, + _ => LIST_SIZE, + })), + ), + ]; + + for (name, list) in inputs { + let (expression, batch) = reduce_expression(&list); + c.bench_with_input( + BenchmarkId::new("array_reduce", name), + &batch, + |b, batch| b.iter(|| expression.evaluate(batch).unwrap()), + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions-nested/src/array_reduce.rs b/datafusion/functions-nested/src/array_reduce.rs new file mode 100644 index 0000000000000..9a26d1a6c1c80 --- /dev/null +++ b/datafusion/functions-nested/src/array_reduce.rs @@ -0,0 +1,340 @@ +// 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. + +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_reduce function. + +use std::sync::Arc; + +use arrow::{ + array::{Array, ArrayRef, AsArray, BooleanArray, UInt64Array, new_null_array}, + compute::{kernels::zip::zip, take, take_arrays}, + datatypes::{DataType, Field, FieldRef}, +}; +use datafusion_common::{ + Result, exec_err, internal_datafusion_err, plan_err, + utils::{adjust_offsets_for_slice, list_values}, +}; +use datafusion_expr::{ + ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, + Volatility, +}; +use datafusion_macros::user_doc; + +make_higher_order_function_expr_and_func!( + ArrayReduce, + array_reduce, + array initial merge, + "reduces an array to a single value using a binary lambda", + array_reduce_higher_order_function +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Reduces an array to a single value by applying a binary lambda to the accumulator and each array element in order.", + syntax_example = "array_reduce(array, initial, (acc, value) -> expression)", + sql_example = r#"```sql +> select array_reduce([1, 2, 3], 0, (acc, value) -> acc + value); ++----------------------------------------------------------------+ +| array_reduce([1,2,3],0,(acc,value) -> acc + value) | ++----------------------------------------------------------------+ +| 6 | ++----------------------------------------------------------------+ +```"#, + argument(name = "array", description = "Array expression to reduce."), + argument(name = "initial", description = "Initial accumulator value."), + argument( + name = "merge", + description = "Binary lambda whose parameters are the accumulator and current array element." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayReduce { + signature: HigherOrderSignature, +} + +impl Default for ArrayReduce { + fn default() -> Self { + Self::new() + } +} + +impl ArrayReduce { + pub fn new() -> Self { + Self { + signature: HigherOrderSignature::exact( + vec![ + ValueOrLambda::Value(()), + ValueOrLambda::Value(()), + ValueOrLambda::Lambda(()), + ], + Volatility::Immutable, + ), + } + } +} + +impl HigherOrderUDFImpl for ArrayReduce { + fn name(&self) -> &str { + "array_reduce" + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn coerce_value_types(&self, arg_types: &[DataType]) -> Result> { + let [list, initial] = arg_types else { + return plan_err!( + "{} requires two value arguments, got {}", + self.name(), + arg_types.len() + ); + }; + + let list = match list { + DataType::List(_) | DataType::LargeList(_) => list.clone(), + DataType::ListView(field) | DataType::FixedSizeList(field, _) => { + DataType::List(Arc::clone(field)) + } + DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)), + DataType::Null => DataType::new_list(DataType::Null, true), + other => { + return plan_err!( + "{} expected a list as first argument, got {other}", + self.name() + ); + } + }; + + Ok(vec![list, initial.clone()]) + } + + fn lambda_parameters( + &self, + step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + let [ + ValueOrLambda::Value(list), + ValueOrLambda::Value(initial), + ValueOrLambda::Lambda(merge), + ] = fields + else { + return plan_err!( + "{} expects an array, an initial value, and a lambda", + self.name() + ); + }; + + let element = match list.data_type() { + DataType::List(field) | DataType::LargeList(field) => Arc::clone(field), + other => { + return plan_err!( + "{} expected a list as first argument, got {other}", + self.name() + ); + } + }; + + match (step, merge) { + (0, None) => Ok(LambdaParametersProgress::Partial(vec![Some(vec![ + Arc::clone(initial), + element, + ])])), + (0 | 1, Some(accumulator)) => Ok(LambdaParametersProgress::Complete(vec![ + vec![Arc::clone(accumulator), element], + ])), + _ => Err(internal_datafusion_err!( + "{} could not resolve its accumulator type at step {step}", + self.name() + )), + } + } + + fn coerce_values_for_lambdas( + &self, + fields: &[ValueOrLambda], + ) -> Result>> { + let [ + ValueOrLambda::Value(list), + ValueOrLambda::Value(_initial), + ValueOrLambda::Lambda(merge), + ] = fields + else { + return plan_err!( + "{} expects an array, an initial value, and a lambda", + self.name() + ); + }; + + Ok(Some(vec![list.clone(), merge.clone()])) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let [ + ValueOrLambda::Value(list), + ValueOrLambda::Value(initial), + ValueOrLambda::Lambda(merge), + ] = args.arg_fields + else { + return plan_err!( + "{} expects an array, an initial value, and a lambda", + self.name() + ); + }; + + Ok(Arc::new(Field::new( + "", + merge.data_type().clone(), + list.is_nullable() || initial.is_nullable() || merge.is_nullable(), + ))) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + let [list, initial, merge] = args.args.as_slice() else { + return exec_err!( + "{} expects an array, an initial value, and a lambda", + self.name() + ); + }; + let ( + ValueOrLambda::Value(list), + ValueOrLambda::Value(initial), + ValueOrLambda::Lambda(merge), + ) = (list, initial, merge) + else { + return exec_err!( + "{} expects an array, an initial value, and a lambda", + self.name() + ); + }; + + let list = list.to_array(args.number_rows)?; + let values = list_values(list.as_ref())?; + let offsets: Vec = match list.data_type() { + DataType::List(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|offset| usize::try_from(*offset)) + .collect::>() + .map_err(|error| { + internal_datafusion_err!("invalid list offset: {error}") + })?, + DataType::LargeList(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|offset| usize::try_from(*offset)) + .collect::>() + .map_err(|error| { + internal_datafusion_err!("invalid list offset: {error}") + })?, + other => return exec_err!("{} expected a list, got {other}", self.name()), + }; + + let mut accumulator = initial.clone().to_array(args.number_rows)?; + if list.null_count() > 0 { + let valid_lists = BooleanArray::from( + (0..list.len()) + .map(|row| list.is_valid(row)) + .collect::>(), + ); + let null_accumulator = + new_null_array(accumulator.data_type(), accumulator.len()); + accumulator = zip(&valid_lists, &accumulator, &null_accumulator)?; + } + + let max_len = offsets + .windows(2) + .map(|pair| pair[1] - pair[0]) + .max() + .unwrap_or(0); + + for position in 0..max_len { + let mut source_indices = Vec::with_capacity(list.len()); + let mut row_indices = Vec::with_capacity(list.len()); + let mut scatter_indices = Vec::with_capacity(list.len()); + + for row in 0..list.len() { + let active = accumulator.is_valid(row) + && list.is_valid(row) + && position < offsets[row + 1] - offsets[row]; + if active { + source_indices.push(u64::try_from(offsets[row] + position).map_err( + |error| internal_datafusion_err!("invalid list index: {error}"), + )?); + row_indices.push(u64::try_from(row).map_err(|error| { + internal_datafusion_err!("invalid row index: {error}") + })?); + scatter_indices.push(Some( + u64::try_from(source_indices.len() - 1).map_err(|error| { + internal_datafusion_err!("invalid scatter index: {error}") + })?, + )); + } else { + scatter_indices.push(None); + } + } + + if source_indices.is_empty() { + break; + } + + if source_indices.len() == list.len() { + let elements = + take(values.as_ref(), &UInt64Array::from(source_indices), None)?; + let accumulator_param = || Ok(Arc::clone(&accumulator)); + let element_param = || Ok(Arc::clone(&elements)); + accumulator = merge + .evaluate(&[&accumulator_param, &element_param], |arrays| { + Ok(arrays.to_vec()) + })? + .into_array(list.len())?; + continue; + } + + let row_indices: ArrayRef = Arc::new(UInt64Array::from(row_indices)); + let active_accumulator = take(accumulator.as_ref(), &row_indices, None)?; + let elements = + take(values.as_ref(), &UInt64Array::from(source_indices), None)?; + let accumulator_param = || Ok(Arc::clone(&active_accumulator)); + let element_param = || Ok(Arc::clone(&elements)); + let merged = merge + .evaluate(&[&accumulator_param, &element_param], |arrays| { + Ok(take_arrays(arrays, &row_indices, None)?) + })? + .into_array(row_indices.len())?; + + let scatter_indices = UInt64Array::from(scatter_indices); + let active_mask = BooleanArray::from( + scatter_indices + .iter() + .map(|index| index.is_some()) + .collect::>(), + ); + let expanded = take(merged.as_ref(), &scatter_indices, None)?; + accumulator = zip(&active_mask, &expanded, &accumulator)?; + } + + Ok(ColumnarValue::Array(accumulator)) + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs index 2c7bd25d7dbcd..f0693da2fd7be 100644 --- a/datafusion/functions-nested/src/lib.rs +++ b/datafusion/functions-nested/src/lib.rs @@ -49,6 +49,7 @@ pub mod array_first; pub mod array_has; pub mod array_normalize; pub mod array_product; +pub mod array_reduce; pub mod array_scale; pub mod array_subtract; pub mod array_sum; @@ -106,6 +107,7 @@ pub mod expr_fn { pub use super::array_has::array_has_any; pub use super::array_normalize::array_normalize; pub use super::array_product::array_product; + pub use super::array_reduce::array_reduce; pub use super::array_scale::array_scale; pub use super::array_subtract::array_subtract; pub use super::array_sum::array_sum; @@ -225,6 +227,7 @@ pub fn all_default_higher_order_functions() -> Vec> { array_any_match::array_any_match_higher_order_function(), array_filter::array_filter_higher_order_function(), array_first::array_first_higher_order_function(), + array_reduce::array_reduce_higher_order_function(), array_transform::array_transform_higher_order_function(), ] } diff --git a/datafusion/sqllogictest/test_files/array/array_reduce.slt b/datafusion/sqllogictest/test_files/array/array_reduce.slt new file mode 100644 index 0000000000000..23f8072eeb190 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array/array_reduce.slt @@ -0,0 +1,98 @@ +# 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_reduce Tests +############# + +statement ok +set datafusion.sql_parser.dialect = databricks; + +query I +SELECT array_reduce([1, 2, 3], 0, (acc, value) -> acc + value); +---- +6 + +query I +SELECT array_reduce([1, 2, 3], 0, (acc, value) -> acc + value * value); +---- +14 + +query T +SELECT array_reduce(['a', 'b', 'c'], '', (acc, value) -> acc || value); +---- +abc + +query I +SELECT array_reduce([], 10, (acc, value) -> acc + value); +---- +10 + +query I +SELECT array_reduce(NULL, 10, (acc, value) -> acc + value); +---- +NULL + +query I +SELECT array_reduce([1, 2, 3], NULL::BIGINT, (acc, value) -> acc + value); +---- +NULL + +query I +SELECT array_reduce([1, NULL, 2], 0, (acc, value) -> acc + value); +---- +NULL + +query I +SELECT array_reduce([1, NULL, 2], 0, (acc, value) -> acc + coalesce(value, 0)); +---- +3 + +# A null merge result is terminal and cannot be recovered by a later element. +query I +SELECT array_reduce( + [1, 2], + 0, + (acc, value) -> CASE WHEN value = 1 THEN NULL ELSE coalesce(acc, 0) + value END +); +---- +NULL + +query R +SELECT array_reduce([1.2, 2.3], 0, (acc, value) -> acc + value); +---- +3.5 + +statement ok +CREATE TABLE reduce_t (values ARRAY, initial BIGINT, extra BIGINT) +AS VALUES +([1, 2, 3], 0, 10), +([4], 5, 20), +([], 7, 30), +(NULL, 8, 40); + +query I rowsort +SELECT array_reduce(values, initial, (acc, value) -> acc + value + extra) +FROM reduce_t; +---- +29 +36 +7 +NULL + +statement ok +set datafusion.sql_parser.dialect = generic;