perf: Compute array_length of first dimension from offsets

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"

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);

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(())
+    }
+}