perf: compute cardinality from list/map offsets

datafusion/functions-nested/Cargo.toml

@@ -137,3 +137,7 @@ name = "array_resize"
 [[bench]]
 harness = false
 name = "array_range"
+
+[[bench]]
+harness = false
+name = "cardinality"

datafusion/functions-nested/benches/cardinality.rs

@@ -0,0 +1,119 @@
+// 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::{
+    Array, ArrayRef, GenericListArray, Int32Array, MapArray, StructArray,
+};
+use arrow::buffer::{NullBuffer, 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::cardinality::cardinality_udf;
+use std::hint::black_box;
+use std::sync::Arc;
+
+fn list_array<O: arrow::array::OffsetSizeTrait>(
+    values: ArrayRef,
+    rows: usize,
+    width: usize,
+    nulls: Option<NullBuffer>,
+) -> ArrayRef {
+    Arc::new(GenericListArray::<O>::new(
+        Arc::new(Field::new_list_field(values.data_type().clone(), true)),
+        OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)),
+        values,
+        nulls,
+    ))
+}
+
+fn bench_cardinality(c: &mut Criterion) {
+    let mut group = c.benchmark_group("cardinality");
+    let udf = cardinality_udf();
+    let return_field = Arc::new(Field::new("cardinality", DataType::UInt64, true));
+    let config_options = Arc::new(ConfigOptions::default());
+
+    let rows = 8192;
+    let width = 32;
+    let values = Arc::new(Int32Array::from_iter_values(
+        (0..rows * width).map(|i| i as i32),
+    )) as ArrayRef;
+    let flat = list_array::<i32>(Arc::clone(&values), rows, width, None);
+    let large = list_array::<i64>(
+        Arc::clone(&values),
+        rows,
+        width,
+        Some(NullBuffer::from(
+            (0..rows).map(|row| row % 5 != 0).collect::<Vec<_>>(),
+        )),
+    );
+    let entries = StructArray::from(vec![
+        (
+            Arc::new(Field::new("key", DataType::Int32, false)),
+            Arc::clone(&values),
+        ),
+        (
+            Arc::new(Field::new("value", DataType::Int32, true)),
+            Arc::clone(&values),
+        ),
+    ]);
+    let map = Arc::new(MapArray::new(
+        Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+        OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows)),
+        entries,
+        None,
+        false,
+    )) as ArrayRef;
+    // Nested lists exercise recursive cardinality: four lists of eight elements.
+    let children = list_array::<i32>(values, rows * 4, 8, None);
+    let nested = list_array::<i32>(children, rows, 4, None);
+
+    for (name, array) in [
+        ("list/valid", Arc::clone(&flat)),
+        ("large_list/nullable", large),
+        ("map/valid", map),
+        ("list/valid", flat.slice(0, 1)),
+        ("nested_list/valid", nested),
+    ] {
+        let number_rows = array.len();
+        let id = BenchmarkId::new(name, format!("{number_rows}x{width}"));
+        let arg_fields = vec![Arc::new(Field::new(
+            "array",
+            array.data_type().clone(),
+            true,
+        ))];
+        let input = ColumnarValue::Array(array);
+        group.bench_function(id, |b| {
+            b.iter(|| {
+                black_box(
+                    udf.invoke_with_args(ScalarFunctionArgs {
+                        args: vec![input.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_cardinality);
+criterion_main!(benches);

datafusion/functions-nested/src/cardinality.rs

@@ -18,9 +18,8 @@
 //! [`ScalarUDFImpl`] definitions for cardinality function.
 
 use crate::utils::make_scalar_function;
-use arrow::array::{
-    Array, ArrayRef, GenericListArray, MapArray, OffsetSizeTrait, UInt64Array,
-};
+use arrow::array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait, UInt64Array};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::datatypes::{
     DataType,
     DataType::{
@@ -129,25 +128,34 @@ fn cardinality_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
         }
         Map(_, _) => {
             let map_array = as_map_array(array)?;
-            generic_map_cardinality(map_array)
+            Ok(cardinality_from_offsets(
+                map_array.offsets(),
+                map_array.nulls(),
+            ))
         }
         arg_type => {
             exec_err!("cardinality does not support type {arg_type}")
         }
     }
 }
 
-fn generic_map_cardinality(array: &MapArray) -> Result<ArrayRef> {
-    let result: UInt64Array = array
-        .iter()
-        .map(|opt_arr| opt_arr.map(|arr| arr.len() as u64))
-        .collect();
-    Ok(Arc::new(result))
+fn cardinality_from_offsets<O: OffsetSizeTrait>(
+    offsets: &OffsetBuffer<O>,
+    nulls: Option<&NullBuffer>,
+) -> ArrayRef {
+    let values = offsets.lengths().map(|len| len as u64).collect::<Vec<_>>();
+    Arc::new(UInt64Array::new(values.into(), nulls.cloned()))
 }
 
 fn generic_list_cardinality<O: OffsetSizeTrait>(
     array: &GenericListArray<O>,
 ) -> Result<ArrayRef> {
+    // Nested lists require recursive counting; for all other lists, we can
+    // compute the cardinality from offsets, which is much faster.
+    if !array.values().data_type().is_list() {
+        return Ok(cardinality_from_offsets(array.offsets(), array.nulls()));
+    }
+
     let result = array
         .iter()
         .map(|arr| match arr {
@@ -198,3 +206,70 @@ where
         })
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Int32Array, MapArray, StructArray};
+    use arrow::datatypes::Field;
+
+    fn check_slices(array: &dyn Array) -> Result<()> {
+        let expected = UInt64Array::from(vec![Some(1), Some(2), Some(0), None]);
+        // Slices retain nonzero offsets into the values and validity buffers.
+        for (offset, len) in [(0, 4), (1, 3), (2, 0)] {
+            let result = cardinality_inner(&[array.slice(offset, len)])?;
+            assert_eq!(result.as_ref(), &expected.slice(offset, len));
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn cardinality_flat_list_offsets() -> Result<()> {
+        fn check<O: OffsetSizeTrait>() -> Result<()> {
+            let values = Arc::new(Int32Array::from(vec![
+                Some(1),
+                None,
+                Some(3),
+                Some(4),
+                Some(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)
+        }
+        check::<i32>()?;
+        check::<i64>()
+    }
+
+    #[test]
+    fn cardinality_map_offsets() -> Result<()> {
+        let entries = StructArray::from(vec![
+            (
+                Arc::new(Field::new("key", DataType::Int32, false)),
+                Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as ArrayRef,
+            ),
+            (
+                Arc::new(Field::new("value", DataType::Int32, true)),
+                Arc::new(Int32Array::from(vec![
+                    Some(1),
+                    None,
+                    Some(3),
+                    Some(4),
+                    Some(5),
+                ])) as ArrayRef,
+            ),
+        ]);
+        let array = MapArray::new(
+            Arc::new(Field::new("entries", entries.data_type().clone(), false)),
+            OffsetBuffer::from_lengths([1, 2, 0, 2]),
+            entries,
+            Some(NullBuffer::from(vec![true, true, true, false])),
+            false,
+        );
+        check_slices(&array)
+    }
+}