diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index ed5a89b8e3e72..cfa642f6205da 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 0000000000000..ba6f9abdbd339 --- /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 0f5055a5f9598..a5f051b8d83c0 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 { - 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::>()?; - - Ok(Arc::new(result) as ArrayRef) - }}; -} - fn array_length_inner(args: &[ArrayRef]) -> Result { - 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::(args), - LargeList(_) => general_array_length::(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 { + let lengths: Vec = 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 { - 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 { + 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(array: &[ArrayRef]) -> Result { - array_length_impl!(as_generic_list_array::(&array[0])?, array.get(1)) +fn lengths_at_dimension( + rows: impl Iterator>, + dimension: &Int64Array, +) -> Result { + let result = rows + .zip(dimension.iter()) + .map(|(row, dim)| compute_array_length(row, dim)) + .collect::>()?; + 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() -> Result<()> { + let values = Arc::new(Int32Array::new_null(5)); + let array = GenericListArray::::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::()?; + check::() + } + + #[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(()) + } +}