This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25121-40488988ad596c9b093ad60e1453430d803ce33c in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4086969669f96c0d5de03f442f2dba1869d2ee80 Author: Neil Conway <[email protected]> AuthorDate: Wed Sep 9 23:06:06 2026 +0000 perf: Compute `array_length` of first dimension from offsets (#25121) ## Which issue does this PR close? - Closes #25119. ## Rationale for this change `array_length` worked by forming a slice of values for each logical row in the given dimension, and then appending the length of that slice to the result. That is inefficient; when computing the length of the first array dimension, we can get that from the offset buffer instead, which is ~200x faster for typical inputs. For other array dimensions, we fallback to the previous code path. Benchmarks: (M4 Max) - list/8192, 175.379 µs -> 1.137 µs, -99.35% - fixed_size_list/8192, 174.770 µs -> 0.963 µs, -99.45% - list/1, 0.171 µs -> 0.088 µs, -48.48% ## What changes are included in this PR? * Compute `array_length` of dimension 1 for `List` and `LargeList` from offsets * Compute `array_length` of dimension 1 for `FixedSizeList` from the FSL's size * Treat `array_length(x, 1)` as equivalent to `array_length(x)` * Simplify the code to remove a no-longer-necessary macro * Add unit tests * Add benchmark for `array_length` ## What is the testing strategy for this PR? Existing tests pass; new unit tests added. ## Are there any user-facing changes? No. --- datafusion/functions-nested/Cargo.toml | 4 + .../functions-nested/benches/array_length.rs | 78 +++++++++ datafusion/functions-nested/src/length.rs | 176 +++++++++++++++++---- 3 files changed, 224 insertions(+), 34 deletions(-) diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index ed5a89b8e3..cfa642f620 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -74,6 +74,10 @@ datafusion-physical-expr = { workspace = true } harness = false name = "array_concat" +[[bench]] +harness = false +name = "array_length" + [[bench]] harness = false name = "array_min_max" diff --git a/datafusion/functions-nested/benches/array_length.rs b/datafusion/functions-nested/benches/array_length.rs new file mode 100644 index 0000000000..ba6f9abdbd --- /dev/null +++ b/datafusion/functions-nested/benches/array_length.rs @@ -0,0 +1,78 @@ +// 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 arrow::array::{ArrayRef, FixedSizeListArray, Int32Array, ListArray}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions_nested::length::array_length_udf; +use std::hint::black_box; +use std::sync::Arc; + +fn bench_array_length(c: &mut Criterion) { + let mut group = c.benchmark_group("array_length"); + let udf = array_length_udf(); + let return_field = Arc::new(Field::new("length", DataType::UInt64, true)); + let config_options = Arc::new(ConfigOptions::default()); + + let rows = 8192; + let width = 32; + let values = Arc::new(Int32Array::new_null(rows * width)) as ArrayRef; + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let flat = Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_repeated_length(width, rows), + Arc::clone(&values), + None, + )) as ArrayRef; + let fixed = + Arc::new(FixedSizeListArray::new(field, width as i32, values, None)) as ArrayRef; + + for (name, array) in [ + ("list", Arc::clone(&flat)), + ("fixed_size_list", fixed), + ("list", flat.slice(0, 1)), + ] { + let number_rows = array.len(); + let id = BenchmarkId::new(name, number_rows); + let args = vec![ColumnarValue::Array(array)]; + let arg_fields: Vec<_> = args + .iter() + .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true))) + .collect(); + group.bench_function(id, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_array_length); +criterion_main!(benches); diff --git a/datafusion/functions-nested/src/length.rs b/datafusion/functions-nested/src/length.rs index 0f5055a5f9..a5f051b8d8 100644 --- a/datafusion/functions-nested/src/length.rs +++ b/datafusion/functions-nested/src/length.rs @@ -20,16 +20,16 @@ use crate::utils::make_scalar_function; use arrow::array::{ Array, ArrayRef, FixedSizeListArray, Int64Array, LargeListArray, ListArray, - OffsetSizeTrait, UInt64Array, + UInt64Array, }; use arrow::datatypes::{ DataType, DataType::{FixedSizeList, LargeList, List, UInt64}, }; use datafusion_common::cast::{ - as_fixed_size_list_array, as_generic_list_array, as_int64_array, + as_fixed_size_list_array, as_int64_array, as_large_list_array, as_list_array, }; -use datafusion_common::{Result, exec_err}; +use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, @@ -114,7 +114,13 @@ impl ScalarUDFImpl for ArrayLength { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { - make_scalar_function(array_length_inner)(&args.args) + // An explicit scalar dimension of one uses the same fast path as an + // omitted dimension. + let args = match args.args.as_slice() { + [_, ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))] => &args.args[..1], + args => args, + }; + make_scalar_function(array_length_inner)(args) } fn aliases(&self) -> &[String] { @@ -126,43 +132,65 @@ impl ScalarUDFImpl for ArrayLength { } } -macro_rules! array_length_impl { - ($array:expr, $dimension:expr) => {{ - let array = $array; - let dimension = match $dimension { - Some(d) => as_int64_array(d)?.clone(), - None => Int64Array::from_value(1, array.len()), - }; - let result = array - .iter() - .zip(dimension.iter()) - .map(|(arr, dim)| compute_array_length(arr, dim)) - .collect::<Result<UInt64Array>>()?; - - Ok(Arc::new(result) as ArrayRef) - }}; -} - fn array_length_inner(args: &[ArrayRef]) -> Result<ArrayRef> { - if args.len() != 1 && args.len() != 2 { - return exec_err!("array_length expects one or two arguments"); + match args { + [array] => first_dimension_length(array), + [array, dimension] => nth_dimension_length(array, as_int64_array(dimension)?), + _ => exec_err!("array_length expects one or two arguments"), } +} - match &args[0].data_type() { - List(_) => general_array_length::<i32>(args), - LargeList(_) => general_array_length::<i64>(args), - FixedSizeList(_, _) => fixed_size_array_length(args), - array_type => exec_err!("array_length does not support type '{array_type}'"), - } +/// Returns each row's length along the first dimension. +/// +/// The first dimension counts a row's immediate elements, whatever their +/// type, so the lengths come straight from the offsets or the fixed width +/// without slicing out any row. +fn first_dimension_length(array: &ArrayRef) -> Result<ArrayRef> { + let lengths: Vec<u64> = match array.data_type() { + List(_) => as_list_array(array)? + .offsets() + .lengths() + .map(|len| len as u64) + .collect(), + LargeList(_) => as_large_list_array(array)? + .offsets() + .lengths() + .map(|len| len as u64) + .collect(), + FixedSizeList(_, size) => vec![*size as u64; array.len()], + array_type => { + return exec_err!("array_length does not support type '{array_type}'"); + } + }; + Ok(Arc::new(UInt64Array::new( + lengths.into(), + array.nulls().cloned(), + ))) } -fn fixed_size_array_length(array: &[ArrayRef]) -> Result<ArrayRef> { - array_length_impl!(as_fixed_size_list_array(&array[0])?, array.get(1)) +/// Returns each row's length along the dimension given for that row. +fn nth_dimension_length(array: &ArrayRef, dimension: &Int64Array) -> Result<ArrayRef> { + match array.data_type() { + List(_) => lengths_at_dimension(as_list_array(array)?.iter(), dimension), + LargeList(_) => { + lengths_at_dimension(as_large_list_array(array)?.iter(), dimension) + } + FixedSizeList(..) => { + lengths_at_dimension(as_fixed_size_list_array(array)?.iter(), dimension) + } + array_type => exec_err!("array_length does not support type '{array_type}'"), + } } -/// Dispatch array length computation based on the offset type. -fn general_array_length<O: OffsetSizeTrait>(array: &[ArrayRef]) -> Result<ArrayRef> { - array_length_impl!(as_generic_list_array::<O>(&array[0])?, array.get(1)) +fn lengths_at_dimension( + rows: impl Iterator<Item = Option<ArrayRef>>, + dimension: &Int64Array, +) -> Result<ArrayRef> { + let result = rows + .zip(dimension.iter()) + .map(|(row, dim)| compute_array_length(row, dim)) + .collect::<Result<UInt64Array>>()?; + Ok(Arc::new(result)) } /// Returns the length of a concrete array dimension @@ -207,3 +235,83 @@ fn compute_array_length( } } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{GenericListArray, Int32Array, OffsetSizeTrait}; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + + fn check_slices(array: &dyn Array, expected: &UInt64Array) -> Result<()> { + let udf = ArrayLength::new(); + for (offset, len) in [(0, 4), (1, 3), (2, 0)] { + let array = array.slice(offset, len); + for explicit_dimension in [false, true] { + let mut args = vec![ColumnarValue::Array(Arc::clone(&array))]; + if explicit_dimension { + args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))); + } + let arg_fields = args + .iter() + .map(|arg| Arc::new(Field::new("arg", arg.data_type(), true))) + .collect(); + let result = udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: len, + return_field: Arc::new(Field::new("length", UInt64, true)), + config_options: Arc::new(ConfigOptions::default()), + })?; + let ColumnarValue::Array(result) = result else { + panic!("expected an array result"); + }; + assert_eq!(result.as_ref(), &expected.slice(offset, len)); + } + } + Ok(()) + } + + #[test] + fn array_length_list_offsets() -> Result<()> { + fn check<O: OffsetSizeTrait>() -> Result<()> { + let values = Arc::new(Int32Array::new_null(5)); + let array = GenericListArray::<O>::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + OffsetBuffer::from_lengths([1, 2, 0, 2]), + values, + Some(NullBuffer::from(vec![true, true, true, false])), + ); + check_slices( + &array, + &UInt64Array::from(vec![Some(1), Some(2), Some(0), None]), + ) + } + check::<i32>()?; + check::<i64>() + } + + #[test] + fn array_length_fixed_size_lists() -> Result<()> { + for width in [0, 2] { + let array = FixedSizeListArray::try_new_with_length( + Arc::new(Field::new_list_field(DataType::Int32, true)), + width, + Arc::new(Int32Array::new_null(4 * width as usize)), + Some(NullBuffer::from(vec![true, true, false, true])), + 4, + )?; + check_slices( + &array, + &UInt64Array::from(vec![ + Some(width as u64), + Some(width as u64), + None, + Some(width as u64), + ]), + )?; + } + Ok(()) + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
