diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index ed5a89b8e3e72..80bf8e7238ea0 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -137,3 +137,7 @@ name = "array_resize" [[bench]] harness = false name = "array_range" + +[[bench]] +harness = false +name = "cardinality" diff --git a/datafusion/functions-nested/benches/cardinality.rs b/datafusion/functions-nested/benches/cardinality.rs new file mode 100644 index 0000000000000..ce0148abd5afb --- /dev/null +++ b/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( + values: ArrayRef, + rows: usize, + width: usize, + nulls: Option, +) -> ArrayRef { + Arc::new(GenericListArray::::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::(Arc::clone(&values), rows, width, None); + let large = list_array::( + Arc::clone(&values), + rows, + width, + Some(NullBuffer::from( + (0..rows).map(|row| row % 5 != 0).collect::>(), + )), + ); + 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::(values, rows * 4, 8, None); + let nested = list_array::(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); diff --git a/datafusion/functions-nested/src/cardinality.rs b/datafusion/functions-nested/src/cardinality.rs index 38def2d4e4afe..cd16c1c3592f7 100644 --- a/datafusion/functions-nested/src/cardinality.rs +++ b/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,7 +128,10 @@ fn cardinality_inner(args: &[ArrayRef]) -> Result { } 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}") @@ -137,17 +139,23 @@ fn cardinality_inner(args: &[ArrayRef]) -> Result { } } -fn generic_map_cardinality(array: &MapArray) -> Result { - 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( + offsets: &OffsetBuffer, + nulls: Option<&NullBuffer>, +) -> ArrayRef { + let values = offsets.lengths().map(|len| len as u64).collect::>(); + Arc::new(UInt64Array::new(values.into(), nulls.cloned())) } fn generic_list_cardinality( array: &GenericListArray, ) -> Result { + // 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() -> Result<()> { + let values = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(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) + } + check::()?; + check::() + } + + #[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) + } +}