This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 4086969669 perf: Compute `array_length` of first dimension from
offsets (#25121)
4086969669 is described below
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]
