From c3d5ca84b6382b65b89a0efa698e008fa4acf610 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Fri, 16 Oct 2026 09:52:34 -0400 Subject: [PATCH 1/6] bench: add inner_product, cosine_distance, and array_normalize benchmarks Cover Float64 lists of 4, 32, and 256 elements per row, with the same total number of elements, both without NULL elements and with one NULL element in 10% of rows. --- datafusion/functions-nested/Cargo.toml | 4 + .../functions-nested/benches/vector_math.rs | 116 ++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 datafusion/functions-nested/benches/vector_math.rs diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index 80b94a0e12bfd..443548b8b3fe4 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -157,3 +157,7 @@ name = "cardinality" [[bench]] harness = false name = "range" + +[[bench]] +harness = false +name = "vector_math" diff --git a/datafusion/functions-nested/benches/vector_math.rs b/datafusion/functions-nested/benches/vector_math.rs new file mode 100644 index 0000000000000..8a0cf19637c8c --- /dev/null +++ b/datafusion/functions-nested/benches/vector_math.rs @@ -0,0 +1,116 @@ +// 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. + +//! Benchmarks for `inner_product`, `cosine_distance`, and `array_normalize`, +//! which compute over each row of `Float64` lists. + +use arrow::array::{ArrayRef, Float64Array, ListArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions_nested::array_normalize::ArrayNormalize; +use datafusion_functions_nested::cosine_distance::CosineDistance; +use datafusion_functions_nested::inner_product::InnerProduct; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::hint::black_box; +use std::sync::Arc; + +// Every input has the same number of elements, so narrower rows mean more rows. +const ELEMENTS: usize = 1 << 17; + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + for width in [4, 32, 256] { + // With "null_elements", 10% of rows contain one NULL element, which + // makes the result for that row NULL. + for (nulls, null_rows) in [("no_nulls", 0.0), ("null_elements", 0.1)] { + let left = float_lists(&mut rng, width, null_rows); + let right = float_lists(&mut rng, width, null_rows); + let case = format!("width_{width}/{nulls}"); + bench( + c, + &format!("inner_product/{case}"), + &InnerProduct::new(), + vec![Arc::clone(&left), Arc::clone(&right)], + ); + bench( + c, + &format!("cosine_distance/{case}"), + &CosineDistance::new(), + vec![Arc::clone(&left), right], + ); + bench( + c, + &format!("array_normalize/{case}"), + &ArrayNormalize::new(), + vec![left], + ); + } + } +} + +fn bench(c: &mut Criterion, name: &str, udf: &dyn ScalarUDFImpl, inputs: Vec) { + let arg_types = inputs + .iter() + .map(|input| input.data_type().clone()) + .collect::>(); + let arg_fields = arg_types + .iter() + .enumerate() + .map(|(i, data_type)| { + Arc::new(Field::new(format!("arg_{i}"), data_type.clone(), true)) + }) + .collect(); + let return_type = udf.return_type(&arg_types).unwrap(); + let args = ScalarFunctionArgs { + number_rows: inputs[0].len(), + args: inputs.into_iter().map(ColumnarValue::Array).collect(), + arg_fields, + return_field: Arc::new(Field::new("result", return_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }; + c.bench_function(name, |b| { + b.iter(|| black_box(udf.invoke_with_args(args.clone()).unwrap())) + }); +} + +fn float_lists(rng: &mut StdRng, width: usize, null_rows: f64) -> ArrayRef { + let rows = ELEMENTS / width; + let values = (0..ELEMENTS) + .map(|_| rng.random_range(-1.0..1.0)) + .collect::>(); + let nulls = (null_rows > 0.0).then(|| { + let mut valid = vec![true; ELEMENTS]; + for row in 0..rows { + if rng.random_bool(null_rows) { + valid[row * width + rng.random_range(0..width)] = false; + } + } + NullBuffer::from(valid) + }); + Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Float64, true)), + OffsetBuffer::from_repeated_length(width, rows), + Arc::new(Float64Array::new(values.into(), nulls)), + None, + )) +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 8d57847fe7ce35f954a514fadf8f3ce0a15f675d Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Fri, 16 Oct 2026 09:52:34 -0400 Subject: [PATCH 2/6] bench: add Spark array_contains and map_from_entries benchmarks array_contains: lists of 4, 16, and 128 Int32 elements and a fixed-size list of 16, with 10% NULL elements and a value that never occurs, so every row checks whether its list contains a NULL element. map_from_entries: lists of 4, 16, and 64 entries with unique keys, where 10% of rows contain one NULL entry. --- datafusion/spark/Cargo.toml | 8 ++ datafusion/spark/benches/array_contains.rs | 88 ++++++++++++++++++ datafusion/spark/benches/map_from_entries.rs | 98 ++++++++++++++++++++ 3 files changed, 194 insertions(+) create mode 100644 datafusion/spark/benches/array_contains.rs create mode 100644 datafusion/spark/benches/map_from_entries.rs diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 08309a931d96b..52f988a41042e 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -106,3 +106,11 @@ name = "sha2" [[bench]] harness = false name = "floor" + +[[bench]] +harness = false +name = "array_contains" + +[[bench]] +harness = false +name = "map_from_entries" diff --git a/datafusion/spark/benches/array_contains.rs b/datafusion/spark/benches/array_contains.rs new file mode 100644 index 0000000000000..cc2d791f5fbc0 --- /dev/null +++ b/datafusion/spark/benches/array_contains.rs @@ -0,0 +1,88 @@ +// 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. + +//! Benchmarks for Spark's `array_contains`, which returns NULL instead of false +//! when the value is not found and the list contains a NULL element. + +use arrow::array::{ArrayRef, FixedSizeListArray, Int32Array, ListArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_spark::function::array::array_contains::SparkArrayContains; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::hint::black_box; +use std::sync::Arc; + +const ROWS: usize = 8192; + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + // 10% of elements are NULL. The needle never occurs, so every row checks + // whether its list contains a NULL element. + for width in [4, 16, 128] { + let values = int_values(&mut rng, ROWS * width); + let list = ListArray::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + OffsetBuffer::from_repeated_length(width, ROWS), + values, + None, + ); + bench(c, &format!("list/width_{width}"), Arc::new(list)); + } + let width = 16; + let list = FixedSizeListArray::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + width as i32, + int_values(&mut rng, ROWS * width), + None, + ); + bench(c, &format!("fixed_size_list/width_{width}"), Arc::new(list)); +} + +fn int_values(rng: &mut StdRng, len: usize) -> ArrayRef { + let values = (0..len) + .map(|_| rng.random_range(0..100)) + .collect::>(); + let nulls = NullBuffer::from_iter((0..len).map(|_| !rng.random_bool(0.1))); + Arc::new(Int32Array::new(values.into(), Some(nulls))) +} + +fn bench(c: &mut Criterion, name: &str, haystack: ArrayRef) { + let udf = SparkArrayContains::new(); + let args = ScalarFunctionArgs { + number_rows: haystack.len(), + arg_fields: vec![ + Arc::new(Field::new("haystack", haystack.data_type().clone(), true)), + Arc::new(Field::new("needle", DataType::Int32, true)), + ], + args: vec![ + ColumnarValue::Array(haystack), + ColumnarValue::Scalar(ScalarValue::Int32(Some(-1))), + ], + return_field: Arc::new(Field::new("result", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + }; + c.bench_function(&format!("array_contains/{name}"), |b| { + b.iter(|| black_box(udf.invoke_with_args(args.clone()).unwrap())) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/spark/benches/map_from_entries.rs b/datafusion/spark/benches/map_from_entries.rs new file mode 100644 index 0000000000000..95fdce1a0ed13 --- /dev/null +++ b/datafusion/spark/benches/map_from_entries.rs @@ -0,0 +1,98 @@ +// 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. + +//! Benchmarks for Spark's `map_from_entries`, which returns NULL for a row whose +//! list of entries contains a NULL entry. + +use arrow::array::{Array, ArrayRef, Int32Array, ListArray, StructArray}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{DataType, Field, Fields}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, +}; +use datafusion_spark::function::map::map_from_entries::MapFromEntries; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use std::hint::black_box; +use std::sync::Arc; + +const ROWS: usize = 8192; + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + // 10% of rows contain one NULL entry. Keys are unique within each row. + for width in [4, 16, 64] { + let len = ROWS * width; + let mut valid = vec![true; len]; + for row in 0..ROWS { + if rng.random_bool(0.1) { + valid[row * width + rng.random_range(0..width)] = false; + } + } + let entries = StructArray::new( + Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int32, true), + ]), + vec![ + Arc::new(Int32Array::from_iter_values( + (0..len).map(|i| (i % width) as i32), + )), + Arc::new(Int32Array::from_iter_values( + (0..len).map(|_| rng.random_range(0..100)), + )), + ], + Some(NullBuffer::from(valid)), + ); + let list = ListArray::new( + Arc::new(Field::new_list_field(entries.data_type().clone(), true)), + OffsetBuffer::from_repeated_length(width, ROWS), + Arc::new(entries), + None, + ); + bench(c, &format!("width_{width}"), Arc::new(list)); + } +} + +fn bench(c: &mut Criterion, name: &str, entries: ArrayRef) { + let udf = MapFromEntries::new(); + let arg_fields = vec![Arc::new(Field::new( + "entries", + entries.data_type().clone(), + true, + ))]; + let return_field = udf + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &[None], + }) + .unwrap(); + let args = ScalarFunctionArgs { + number_rows: entries.len(), + args: vec![ColumnarValue::Array(entries)], + arg_fields, + return_field, + config_options: Arc::new(ConfigOptions::default()), + }; + c.bench_function(&format!("map_from_entries/{name}"), |b| { + b.iter(|| black_box(udf.invoke_with_args(args.clone()).unwrap())) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From c5615fe885da4a0efe0177a62e6ac5b6cb24d360 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Fri, 16 Oct 2026 09:53:39 -0400 Subject: [PATCH 3/6] refactor: move SetBitCounter to datafusion-common and add rows_have_nulls Make SetBitCounter public in datafusion_common::utils so other crates can count the set bits in each row of a list in one pass, and add count_range for the set bits within one row, which prepare_list_filter now uses for the elements stored under NULL rows. Add rows_have_nulls, which returns a bitmap of the list rows that contain a NULL element. It counts from the list's first element, which can be far into the validity bitmap when the list is sliced, builds the bitmap with BooleanBuffer::collect_bool, and reuses each row's end count as the next row's start. --- datafusion/common/src/utils/mod.rs | 135 ++++++++++++++++++++++- datafusion/functions-nested/src/utils.rs | 59 +--------- 2 files changed, 136 insertions(+), 58 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 6727ab7dcce80..c92884db1be52 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -38,7 +38,7 @@ use arrow::array::{ MutableArrayData, PrimitiveArray, layout, make_array, }; use arrow::array::{LargeListViewArray, ListViewArray}; -use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::kernels::cmp::eq; use arrow::compute::kernels::length::length; use arrow::compute::{SortColumn, SortOptions, nullif, partition}; @@ -1266,6 +1266,94 @@ pub fn offset_span(offsets: &OffsetBuffer) -> (usize, usi (start, offsets.last().as_usize() - start) } +/// Counts the set bits that precede a sequence of non-decreasing positions in a +/// bitmap, reading each 64-bit word of the bitmap once. +/// +/// This suits a single pass over the rows of a list, where each call costs a +/// mask and a popcount when the position is in the current word. `words` yields +/// the bitmap's 64-bit words, starting with its first bit, such as +/// [`BooleanBuffer::bit_chunks`] with [`BitChunks::iter_padded`]. +/// +/// [`BooleanBuffer::bit_chunks`]: arrow::buffer::BooleanBuffer::bit_chunks +/// [`BitChunks::iter_padded`]: arrow::util::bit_chunk_iterator::BitChunks::iter_padded +/// +/// ``` +/// # use arrow::buffer::BooleanBuffer; +/// # use datafusion_common::utils::SetBitCounter; +/// let bits = BooleanBuffer::from(vec![true, false, true, true, false]); +/// let mut counter = SetBitCounter::new(bits.bit_chunks().iter_padded()); +/// assert_eq!(counter.count_before(1), 1); +/// assert_eq!(counter.count_range(1, 5), 2); +/// ``` +pub struct SetBitCounter { + /// The bitmap's remaining words, with the final word padded with zeros. + words: I, + /// The word containing the most recent position. + word: u64, + /// Position of the first bit of `word`. + word_start: usize, + /// Set bits before `word_start`. + count_before_word: usize, +} + +impl> SetBitCounter { + pub fn new(mut words: I) -> Self { + let word = words.next().unwrap_or(0); + Self { + words, + word, + word_start: 0, + count_before_word: 0, + } + } + + /// Returns the number of set bits before `end`. `end` must be no less than + /// in the previous call, and no greater than the bitmap's length. + #[inline] + pub fn count_before(&mut self, end: usize) -> usize { + debug_assert!(end >= self.word_start); + while end - self.word_start >= 64 { + self.count_before_word += self.word.count_ones() as usize; + self.word = self.words.next().unwrap_or(0); + self.word_start += 64; + } + let mask = (1u64 << (end - self.word_start)) - 1; + self.count_before_word + (self.word & mask).count_ones() as usize + } + + /// Returns the number of set bits in `start..end`. `start` must be no less + /// than the position in the previous call. + #[inline] + pub fn count_range(&mut self, start: usize, end: usize) -> usize { + let before = self.count_before(start); + self.count_before(end) - before + } +} + +/// Returns a bitmap whose bit `i` is set if row `i` of a list contains a NULL +/// element, given the elements' validity `nulls`. The rows are consecutive: the +/// first starts at element `start`, and each ends at the next value of `ends`, +/// such as the offsets of a list after its first. +pub fn rows_have_nulls( + nulls: &NullBuffer, + start: usize, + mut ends: impl ExactSizeIterator, +) -> BooleanBuffer { + // Count from `start`, which can be far into `nulls` when the list is sliced. + let visible = nulls.inner().slice(start, nulls.len() - start); + let mut valid = SetBitCounter::new(visible.bit_chunks().iter_padded()); + let mut row_start = start; + let mut valid_before_row = 0; + // `collect_bool` visits the rows in order, as the counter requires. + BooleanBuffer::collect_bool(ends.len(), |_| { + let row_end = ends.next().unwrap_or(row_start); + let valid_through_row = valid.count_before(row_end - start); + let has_null = valid_through_row - valid_before_row < row_end - row_start; + (row_start, valid_before_row) = (row_end, valid_through_row); + has_null + }) +} + /// Returns the inner values of a list, or an error otherwise /// For [`ListArray`] and [`LargeListArray`], if it's sliced, it returns a /// sliced array too. Therefore, too reconstruct a list using it, @@ -1627,6 +1715,51 @@ mod tests { #[cfg(feature = "sql")] use sqlparser::ast::Ident; + #[test] + fn rows_have_nulls_matches_per_row_null_counts() { + // Rows of 0 to 6 elements where every 7th element is NULL, read from + // element 24 on, partway through the first 64-bit word. + let nulls = NullBuffer::from_iter((0..400).map(|i| i % 7 != 3)); + let offsets = OffsetBuffer::::from_lengths((0..90).map(|row| row % 7)); + let offsets = offsets.slice(10, 70); + let start = offsets[0] as usize; + assert_eq!(start, 24); + let has_nulls = + rows_have_nulls(&nulls, start, offsets[1..].iter().map(|&end| end as usize)); + let expected = offsets + .windows(2) + .map(|w| { + nulls + .slice(w[0] as usize, (w[1] - w[0]) as usize) + .null_count() + != 0 + }) + .collect::>(); + assert_eq!(has_nulls.iter().collect::>(), expected); + } + + #[test] + fn set_bit_counter_matches_per_position_counts() { + // A sliced bitmap whose length is a whole number of 64-bit words. + let bits = + BooleanBuffer::collect_bool(200, |i| i % 3 != 0 && i % 7 != 0).slice(5, 192); + let mut counter = SetBitCounter::new(bits.bit_chunks().iter_padded()); + let mut expected = 0; + for end in 0..=bits.len() { + assert_eq!(counter.count_before(end), expected, "end {end}"); + if end < bits.len() { + expected += usize::from(bits.value(end)); + } + } + + // Ranges that skip some positions and span several words. + let mut counter = SetBitCounter::new(bits.bit_chunks().iter_padded()); + for (start, end) in [(0, 0), (3, 10), (10, 70), (100, 100), (130, 192)] { + let expected = (start..end).filter(|&i| bits.value(i)).count(); + assert_eq!(counter.count_range(start, end), expected, "{start}..{end}"); + } + } + #[test] fn test_list_scalar_cast_preserves_target_element_metadata() -> Result<()> { let input = diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index 73bf3cbf4c84f..5eaab0c99b27e 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -30,7 +30,7 @@ use datafusion_common::cast::{ as_fixed_size_list_array, as_float64_array, as_generic_list_array, as_large_list_array, as_large_list_view_array, as_list_array, as_list_view_array, }; -use datafusion_common::utils::{offset_span, offset_span_len}; +use datafusion_common::utils::{SetBitCounter, offset_span, offset_span_len}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err}; use datafusion_expr::ColumnarValue; @@ -95,8 +95,7 @@ pub(crate) fn prepare_list_filter( .append_packed_range(bit_start..bit_start + len, keep.values()); } } else if len != 0 { - let set_before_row = set_bits.count_before(start); - hidden_count += set_bits.count_before(end) - set_before_row; + hidden_count += set_bits.count_range(start, end); keep_mask .get_or_insert_with(|| { let mut builder = BooleanBufferBuilder::new(keep.len()); @@ -120,45 +119,6 @@ pub(crate) fn prepare_list_filter( (offsets, keep_mask) } -/// Counts the set bits that precede a sequence of non-decreasing positions in a -/// bitmap, reading each 64-bit word of the bitmap once. -struct SetBitCounter { - /// The bitmap's remaining words, with the final word padded with zeros. - words: I, - /// The word containing the most recent position. - word: u64, - /// Position of the first bit of `word`. - word_start: usize, - /// Set bits before `word_start`. - count_before_word: usize, -} - -impl> SetBitCounter { - fn new(mut words: I) -> Self { - let word = words.next().unwrap_or(0); - Self { - words, - word, - word_start: 0, - count_before_word: 0, - } - } - - /// Returns the number of set bits before `end`. `end` must be no less than - /// in the previous call, and no greater than the bitmap's length. - #[inline] - fn count_before(&mut self, end: usize) -> usize { - debug_assert!(end >= self.word_start); - while end - self.word_start >= 64 { - self.count_before_word += self.word.count_ones() as usize; - self.word = self.words.next().unwrap_or(0); - self.word_start += 64; - } - let mask = (1u64 << (end - self.word_start)) - 1; - self.count_before_word + (self.word & mask).count_ones() as usize - } -} - /// Computes the return type of a function that produces a list with the same /// inner field as `array_type`, plus an element that may be null when /// `element_nullable` is set. @@ -758,21 +718,6 @@ pub(crate) mod tests { } } - #[test] - fn set_bit_counter_matches_per_position_counts() { - // A sliced bitmap whose length is a whole number of 64-bit words. - let bits = - BooleanBuffer::collect_bool(200, |i| i % 3 != 0 && i % 7 != 0).slice(5, 192); - let mut counter = SetBitCounter::new(bits.bit_chunks().iter_padded()); - let mut expected = 0; - for end in 0..=bits.len() { - assert_eq!(counter.count_before(end), expected, "end {end}"); - if end < bits.len() { - expected += usize::from(bits.value(end)); - } - } - } - /// Only test internal functions, array-related sql functions will be tested in sqllogictest `array.slt` #[test] fn test_align_array_dimensions() { From 18b2c43a99232651929b36a71f7ece3d036f441a Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sat, 17 Oct 2026 13:26:01 -0400 Subject: [PATCH 4/6] perf: count NULL elements per list row in one pass in nested functions array_remove, inner_product, cosine_distance, and array_normalize sliced a bitmap or the values array for each list row to count its set bits or NULLs. Each slice clones a reference-counted buffer and sets up a new bit count. array_remove now counts each row's matches with SetBitCounter, which reads each 64-bit word of the bitmap once over the whole list, and slices a row's bits only when it has matches to remove. The vector functions get the rows that contain a NULL from rows_have_nulls, through a new rows_with_null_elements helper, and read each row's values as a plain slice of the values buffer. Test inner_product, cosine_distance, and array_remove with a scalar needle through check_sliced_list_behavior, and pad its input by a length that is not a multiple of 64, so the list's elements start partway through a 64-bit word. x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest of 3 alternating rounds; no other commit on the branch changes these functions: no_nulls null_elements inner_product/width_4 -88.8% -81.2% inner_product/width_32 -88.9% -84.0% inner_product/width_256 -79.5% -76.7% cosine_distance/width_4 -78.5% -73.2% cosine_distance/width_32 -57.1% -57.5% cosine_distance/width_256 -17.7% -16.3% array_normalize/width_4 -80.9% -71.0% array_normalize/width_32 -77.2% -73.6% array_normalize/width_256 -65.2% -64.2% array_remove with a scalar needle, lists of 10 elements: array_remove_int64 -11.7% array_remove_n_int64 -14.8% array_remove_all_int64 -9.8% array_remove_strings -3.0% array_remove_boolean -12.7% array_remove_fixed_size_binary -11.6% With lists of 100 and 500 elements, which do more work per row, array_remove moved -2.1% to +0.8%. The nested-list cases, whose list needle takes a path this does not change, moved -1.9% to +2.7%. --- .../functions-nested/src/array_normalize.rs | 18 ++++++----- .../functions-nested/src/cosine_distance.rs | 30 +++++++++++++++---- .../functions-nested/src/inner_product.rs | 28 +++++++++++++---- datafusion/functions-nested/src/remove.rs | 20 +++++++++++-- datafusion/functions-nested/src/utils.rs | 22 ++++++++++++-- 5 files changed, 93 insertions(+), 25 deletions(-) diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs index 386ea5d77bfcd..a1f90edaed3d5 100644 --- a/datafusion/functions-nested/src/array_normalize.rs +++ b/datafusion/functions-nested/src/array_normalize.rs @@ -17,7 +17,9 @@ //! [`ScalarUDFImpl`] definitions for array_normalize function. -use crate::utils::{make_scalar_function, needs_norm_scale, norm_scale}; +use crate::utils::{ + make_scalar_function, needs_norm_scale, norm_scale, rows_with_null_elements, +}; use arrow::array::{ Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, OffsetSizeTrait, }; @@ -149,6 +151,7 @@ fn general_array_normalize(arrays: &[ArrayRef]) -> Result::with_capacity(list_array.len() + 1); new_offsets.push(O::zero()); let mut nulls = NullBufferBuilder::new(list_array.len()); + let null_rows = rows_with_null_elements(list_array); for row in 0..list_array.len() { if list_array.is_null(row) { @@ -161,14 +164,13 @@ fn general_array_normalize(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result| rows.as_ref().is_some_and(|r| r.value(row)); + if has_null(&null_rows1) || has_null(&null_rows2) { builder.append_null(); continue; } - let vals1 = slice1.values(); - let vals2 = slice2.values(); + let vals1 = &values1.values()[start1..end1]; + let vals2 = &values2.values()[start2..end2]; let (mut dot, mut sq1, mut sq2) = dot_and_squares(vals1, vals2, 1.0, 1.0); // Cosine distance does not change when either vector is multiplied by a @@ -255,3 +261,15 @@ fn dot_and_squares( } (dot, sq1, sq2) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + #[test] + fn test_sliced_list() -> datafusion_common::Result<()> { + crate::utils::tests::check_sliced_list_behavior(|input| { + super::cosine_distance_inner(&[Arc::clone(input), Arc::clone(input)]) + }) + } +} diff --git a/datafusion/functions-nested/src/inner_product.rs b/datafusion/functions-nested/src/inner_product.rs index da8ea0a1d0215..278e0b4043686 100644 --- a/datafusion/functions-nested/src/inner_product.rs +++ b/datafusion/functions-nested/src/inner_product.rs @@ -17,8 +17,9 @@ //! [`ScalarUDFImpl`] definitions for inner_product function. -use crate::utils::make_scalar_function; +use crate::utils::{make_scalar_function, rows_with_null_elements}; use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; +use arrow::buffer::BooleanBuffer; use arrow::datatypes::{ DataType, DataType::{FixedSizeList, LargeList, List, Null}, @@ -173,6 +174,9 @@ fn general_inner_product(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result| rows.as_ref().is_some_and(|r| r.value(row)); + if has_null(&null_rows1) || has_null(&null_rows2) { builder.append_null(); continue; } - let vals1 = slice1.values(); - let vals2 = slice2.values(); + let vals1 = &values1.values()[start1..end1]; + let vals2 = &values2.values()[start2..end2]; let mut dot = 0.0; for i in 0..len1 { @@ -212,3 +216,15 @@ fn general_inner_product(arrays: &[ArrayRef]) -> Result datafusion_common::Result<()> { + crate::utils::tests::check_sliced_list_behavior(|input| { + super::inner_product_inner(&[Arc::clone(input), Arc::clone(input)]) + }) + } +} diff --git a/datafusion/functions-nested/src/remove.rs b/datafusion/functions-nested/src/remove.rs index befcc09d22f99..5b94aec9d9826 100644 --- a/datafusion/functions-nested/src/remove.rs +++ b/datafusion/functions-nested/src/remove.rs @@ -25,7 +25,9 @@ use arrow::array::{ use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::cast::as_int64_array; -use datafusion_common::utils::{ListCoercion, offset_span, offset_span_len}; +use datafusion_common::utils::{ + ListCoercion, SetBitCounter, offset_span, offset_span_len, +}; use datafusion_common::{ Result, ScalarValue, exec_err, internal_err, utils::take_function_args, }; @@ -591,6 +593,7 @@ fn general_remove_with_scalar( let needle = scalar_needle.to_array_of_size(1)?; let remove_mask = arrow_ord::cmp::not_distinct(&values_slice, &Scalar::new(needle))?; let remove_bits = remove_mask.values(); + let mut matches = SetBitCounter::new(remove_bits.bit_chunks().iter_padded()); for (row_index, offset_window) in list_offsets.windows(2).enumerate() { if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) { @@ -602,8 +605,7 @@ fn general_remove_with_scalar( let end = offset_window[1].to_usize().unwrap() - first_offset; let row_len = end - start; - let row_remove_bits = remove_bits.slice(start, row_len); - let num_to_remove = row_remove_bits.count_set_bits(); + let num_to_remove = matches.count_range(start, end); if num_to_remove == 0 { mutable.try_extend(0, start, end)?; @@ -612,6 +614,7 @@ fn general_remove_with_scalar( } let removals_to_apply = max_removals.min(num_to_remove as i64) as usize; + let row_remove_bits = remove_bits.slice(start, row_len); // Iterate only over the removal positions via set_indices. This is // efficient when the number of removals is small relative to the row @@ -1176,4 +1179,15 @@ mod tests { super::array_remove_internal(input, &element, &vec![Some(1); input.len()]) }) } + + #[test] + fn test_sliced_capacity_scalar_needle() -> datafusion_common::Result<()> { + crate::utils::tests::check_sliced_list_behavior(|input| { + super::array_remove_with_scalar_args( + input, + &ScalarValue::Float64(Some(4.0)), + 1, + ) + }) + } } diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index 5eaab0c99b27e..4b8c5bfbd37da 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -30,7 +30,9 @@ use datafusion_common::cast::{ as_fixed_size_list_array, as_float64_array, as_generic_list_array, as_large_list_array, as_large_list_view_array, as_list_array, as_list_view_array, }; -use datafusion_common::utils::{SetBitCounter, offset_span, offset_span_len}; +use datafusion_common::utils::{ + SetBitCounter, offset_span, offset_span_len, rows_have_nulls, +}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err}; use datafusion_expr::ColumnarValue; @@ -591,6 +593,20 @@ pub(crate) fn needs_norm_scale(sum_of_squares: f64, len: usize) -> bool { !(min_unscaled..f64::INFINITY).contains(&sum_of_squares) } +/// Returns a bitmap whose bit `i` is set if row `i` of `list` contains a NULL +/// element, or `None` if the elements have no validity bitmap. +pub(crate) fn rows_with_null_elements( + list: &GenericListArray, +) -> Option { + let nulls = list.values().nulls()?; + let offsets = list.offsets(); + Some(rows_have_nulls( + nulls, + offsets[0].as_usize(), + offsets[1..].iter().map(|end| end.as_usize()), + )) +} + #[cfg(test)] pub(crate) mod tests { use super::*; @@ -819,7 +835,9 @@ pub(crate) mod tests { pub(crate) fn check_sliced_list_behavior( run: impl Fn(&ArrayRef) -> Result, ) -> Result<()> { - let padding = 8192; + // Not a multiple of 64, so the visible elements start partway through a + // 64-bit word of the child's validity bitmap. + let padding = 8195; let visible = vec![Some(3.0), None, Some(4.0), Some(9.0), Some(9.0)]; let values = Float64Array::from_iter( std::iter::repeat_n(Some(1.0), padding) From 6d14b8be6ea8ea08f27e073ba5b89ce012e9c36e Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sat, 17 Oct 2026 13:26:01 -0400 Subject: [PATCH 5/6] perf: count NULL entries per row in one pass in Spark array_contains and map_from_entries Both sliced the element validity bitmap for each list row to check whether the row contains a NULL. Use rows_have_nulls, which counts the valid elements of each row in one pass over the bitmap. Add tests that a sliced list gives the same rows as the whole list. x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest of 3 alternating rounds; no other commit on the branch changes these functions: array_contains/list/width_4 -67.2% array_contains/list/width_16 -53.6% array_contains/list/width_128 -10.2% array_contains/fixed_size_list/width_16 -48.6% map_from_entries/width_4 -8.2% map_from_entries/width_16 -3.1% map_from_entries/width_64 -2.8% --- .../src/function/array/array_contains.rs | 56 ++++++++++++------ .../src/function/map/map_from_entries.rs | 57 ++++++++++++------- 2 files changed, 77 insertions(+), 36 deletions(-) diff --git a/datafusion/spark/src/function/array/array_contains.rs b/datafusion/spark/src/function/array/array_contains.rs index 5c7cb4be6ff9d..59484dcb365cc 100644 --- a/datafusion/spark/src/function/array/array_contains.rs +++ b/datafusion/spark/src/function/array/array_contains.rs @@ -15,11 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ - Array, AsArray, BooleanArray, BooleanBufferBuilder, GenericListArray, OffsetSizeTrait, -}; +use arrow::array::{Array, AsArray, BooleanArray, GenericListArray, OffsetSizeTrait}; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::DataType; +use datafusion_common::utils::rows_have_nulls; use datafusion_common::{Result, exec_err}; use datafusion_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, @@ -114,13 +113,8 @@ fn compute_row_has_nulls(haystack: &dyn Array) -> Result { let list = haystack.as_fixed_size_list(); let buf = match list.values().nulls() { Some(nulls) => { - let validity = nulls.inner(); let vl = list.value_length() as usize; - let mut builder = BooleanBufferBuilder::new(list.len()); - for i in 0..list.len() { - builder.append(validity.slice(i * vl, vl).count_set_bits() < vl); - } - builder.finish() + rows_have_nulls(nulls, 0, (0..list.len()).map(|row| (row + 1) * vl)) } None => BooleanBuffer::new_unset(list.len()), }; @@ -136,15 +130,12 @@ fn generic_list_row_has_nulls( ) -> Result { let buf = match list.values().nulls() { Some(nulls) => { - let validity = nulls.inner(); let offsets = list.offsets(); - let mut builder = BooleanBufferBuilder::new(list.len()); - for i in 0..list.len() { - let s = offsets[i].as_usize(); - let len = offsets[i + 1].as_usize() - s; - builder.append(validity.slice(s, len).count_set_bits() < len); - } - builder.finish() + rows_have_nulls( + nulls, + offsets[0].as_usize(), + offsets[1..].iter().map(|end| end.as_usize()), + ) } None => BooleanBuffer::new_unset(list.len()), }; @@ -161,3 +152,34 @@ fn mask_with_list_nulls( None => buf, } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, FixedSizeListArray, Int32Array, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Field; + + #[test] + fn row_has_nulls_on_sliced_lists() -> Result<()> { + // Every 7th element and every 6th list row are NULL. + let values = Arc::new(Int32Array::from_iter( + (0..300).map(|i| (i % 7 != 3).then_some(i)), + )); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let list = ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths((0..100).map(|row| row % 5)), + Arc::clone(&values) as ArrayRef, + Some(NullBuffer::from_iter((0..100).map(|row| row % 6 != 0))), + ); + let fixed = FixedSizeListArray::new(field, 3, values, None); + for array in [Arc::new(list) as ArrayRef, Arc::new(fixed)] { + // The first visible element is past the first 64-bit word. + let full = compute_row_has_nulls(array.as_ref())?; + let sliced = compute_row_has_nulls(array.slice(37, 50).as_ref())?; + assert_eq!(sliced, full.slice(37, 50), "{}", array.data_type()); + } + Ok(()) + } +} diff --git a/datafusion/spark/src/function/map/map_from_entries.rs b/datafusion/spark/src/function/map/map_from_entries.rs index 69ce352694bd1..b2ff785d90990 100644 --- a/datafusion/spark/src/function/map/map_from_entries.rs +++ b/datafusion/spark/src/function/map/map_from_entries.rs @@ -21,11 +21,11 @@ use crate::function::map::utils::{ get_list_offsets, get_list_values, map_from_keys_values_offsets_nulls, map_type_from_key_value_types, }; -use arrow::array::{Array, ArrayRef, NullBufferBuilder, StructArray}; +use arrow::array::{Array, ArrayRef, StructArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::config::MapKeyDedupPolicy; -use datafusion_common::utils::take_function_args; +use datafusion_common::utils::{rows_have_nulls, take_function_args}; use datafusion_common::{Result, exec_err, internal_err}; use datafusion_expr::{ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -125,24 +125,14 @@ fn map_from_entries_inner(args: &[ArrayRef], last_value_wins: bool) -> Result Result Result<()> { + // 40 rows of 3 entries with unique keys; every 7th entry is NULL. + let entries = StructArray::new( + Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int32, true), + ]), + vec![ + Arc::new(Int32Array::from_iter_values((0..120).map(|i| i % 3))), + Arc::new(Int32Array::from_iter_values(0..120)), + ], + Some(NullBuffer::from_iter((0..120).map(|i| i % 7 != 3))), + ); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(entries.data_type().clone(), true)), + OffsetBuffer::from_repeated_length(3, 40), + Arc::new(entries), + None, + )); + // The first visible entry is past the first 64-bit word. + let full = map_from_entries_inner(&[Arc::clone(&list)], false)?; + let sliced = map_from_entries_inner(&[list.slice(23, 15)], false)?; + assert_eq!(sliced.as_ref(), full.slice(23, 15).as_ref()); + Ok(()) + } + fn make_entries_field(array_nullable: bool, element_nullable: bool) -> FieldRef { let struct_type = DataType::Struct(Fields::from(vec![ Field::new("key", DataType::Int32, false), From 6225dfc43b1d43c9c599ee994a611875a483f368 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sat, 17 Oct 2026 13:26:01 -0400 Subject: [PATCH 6/6] perf: build array_compact's keep bitmap with prepare_list_filter For lists of non-primitive values, array_compact counted each valid row's non-NULL elements with count_set_bits_offset, which has a fixed cost per call that dominates for short rows, in a loop that otherwise repeated prepare_list_filter. Call prepare_list_filter instead, which counts with SetBitCounter in one pass over the element validity and reuses that validity as the keep bitmap until a NULL row stores elements. The old array_compact loop did no work for NULL rows, but prepare_list_filter counted each NULL row's set bits and appended its cleared bits to the keep bitmap one row at a time. Calling it as it was made array_compact nearly three times slower when most rows were NULL and stored elements. prepare_list_filter now handles the elements between two valid rows together: when a valid row starts after the previous one ended, it counts the set bits in between once and clears them with one append, and it clears any elements after the last valid row at the end. A NULL row only adds its offset. A valid row that directly follows the previous valid row still costs one count, so inputs whose NULL rows store no elements do the same work as before. array_filter also builds its keep bitmap with prepare_list_filter, so it shares this handling. x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest of 3 alternating rounds; no other commit on the branch changes array_compact: dictionary/mixed_nulls/2048 -41.4% dictionary/mixed_nulls/1048576 -40.5% random/utf8/random_sparse -44.4% random/utf8/random -40.1% random/utf8/random_half -39.3% random/utf8/random_dense -65.4% random/utf8_view/random_sparse -56.2% random/utf8_view/random -40.1% random/utf8_view/random_half -47.1% random/utf8_view/random_dense -68.9% types/utf8/mixed -35.7% nulls/utf8/clustered -49.0% types/boolean/mixed -36.0% types/nested/sparse -31.0% types/struct_nested/sparse -28.6% types/fixed_nested/sparse -18.7% types/map/sparse -21.0% types/run/mixed -23.1% strings/utf8_long -19.9% shape/utf8/width_128 -4.2% large_list/utf8 -36.5% child_slice/utf8 -36.2% large/utf8/random_tenth -17.1% large/utf8/random_half -17.2% large/utf8_long/random_half -0.6% binary/random_half -26.9% large_binary/random_half -24.0% strings/null_payload -68.4% With NULL rows that store elements: parent_nulls/utf8/some +0.5% parent_nulls/utf8/mostly +42.5% parent_nulls/nested/mostly +9.5% strings/null_payload_null_rows +0.6% Lists whose rows are mostly NULL but still store elements remain slower than before, because SetBitCounter counts every 64-bit word of the element validity, including the words under NULL rows, which the old loop never read. NULL rows that store elements are rare. Cases with primitive values, which take a separate path, and cases that return before reaching this code moved -2.0% to +5.3%. array_filter's cases with NULL rows moved -5.6% to +8.5%, and its cases without NULL rows, which run none of the code this changes, moved -5.7% to +12.1%. --- .../functions-nested/src/array_compact.rs | 43 +++------------- datafusion/functions-nested/src/utils.rs | 50 +++++++++++++------ 2 files changed, 41 insertions(+), 52 deletions(-) diff --git a/datafusion/functions-nested/src/array_compact.rs b/datafusion/functions-nested/src/array_compact.rs index 12e3638294adf..277f124921dd3 100644 --- a/datafusion/functions-nested/src/array_compact.rs +++ b/datafusion/functions-nested/src/array_compact.rs @@ -17,12 +17,11 @@ //! [`ScalarUDFImpl`] definitions for array_compact function. -use crate::utils::{empty_list_values, make_scalar_function}; +use crate::utils::{empty_list_values, make_scalar_function, prepare_list_filter}; use arrow::array::types::ByteArrayType; use arrow::array::{ - Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, BooleanBufferBuilder, - GenericByteArray, GenericListArray, OffsetSizeTrait, PrimitiveArray, - downcast_primitive_array, + Array, ArrayRef, ArrowPrimitiveType, AsArray, BooleanArray, GenericByteArray, + GenericListArray, OffsetSizeTrait, PrimitiveArray, downcast_primitive_array, }; use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; use arrow::compute::filter; @@ -218,37 +217,9 @@ fn compact_non_primitive( list_array: &GenericListArray, values_nulls: &NullBuffer, ) -> Result<(OffsetBuffer, ArrayRef)> { - let list_offsets = list_array.offsets(); - let first_offset = list_offsets[0].as_usize(); - let mut offsets = Vec::with_capacity(list_array.len() + 1); - offsets.push(O::zero()); - let mut kept_count = 0; - // Element validity is already the keep bitmap unless NULL rows hide elements. - let mut keep_mask = - (list_array.null_count() != 0).then(|| BooleanBufferBuilder::new(values.len())); - for (row, window) in list_offsets.windows(2).enumerate() { - let start = window[0].as_usize() - first_offset; - let len = window[1].as_usize() - window[0].as_usize(); - if list_array.is_valid(row) { - let bit_start = values_nulls.offset() + start; - kept_count += values_nulls.buffer().count_set_bits_offset(bit_start, len); - if let Some(keep_mask) = &mut keep_mask { - // Skip elements belonging to preceding NULL rows. - keep_mask.resize(start); - keep_mask.append_packed_range( - bit_start..bit_start + len, - values_nulls.validity(), - ); - } - } - offsets.push(O::usize_as(kept_count)); - } - let keep_mask = keep_mask - .map(|mut keep_mask| { - keep_mask.resize(values.len()); - keep_mask.finish() - }) - .unwrap_or_else(|| values_nulls.inner().clone()); + // Keep the non-NULL elements, except those stored under NULL rows. + let (offsets, keep_mask) = prepare_list_filter(list_array, values_nulls.inner()); + let kept_count = offsets.last().as_usize(); let output = match values.data_type() { DataType::Utf8 => { copy_byte_spans(values.as_string::(), &keep_mask, kept_count) @@ -264,7 +235,7 @@ fn compact_non_primitive( } _ => filter(values, &BooleanArray::new(keep_mask, None))?, }; - Ok((OffsetBuffer::new(offsets.into()), output)) + Ok((offsets, output)) } /// Copy spans of retained strings or binary values. The keep bitmap excludes diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index 4b8c5bfbd37da..5201ddeb0e3a8 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -81,34 +81,52 @@ pub(crate) fn prepare_list_filter( let mut kept_count = 0; // Set bits stored under NULL rows, which the result does not keep. let mut hidden_count = 0; + // The valid rows so far end at element `covered`, with `set_before_covered` + // set bits before it. Elements between there and the next valid row + // belong to NULL rows, which are handled together when that valid row + // is reached, or after the last row. + let mut covered = 0; + let mut set_before_covered = 0; // NULL list rows are usually empty, so reuse `keep` until a NULL row // stores elements. Copy the bits of the preceding rows when that first // happens. let mut keep_mask: Option = None; - for (row, window) in list_offsets.windows(2).enumerate() { - let start = window[0].as_usize() - first_offset; - let end = window[1].as_usize() - first_offset; - let len = end - start; - let bit_start = keep.offset() + start; - if list.is_valid(row) { - kept_count = set_bits.count_before(end) - hidden_count; - if let Some(keep_mask) = &mut keep_mask { - keep_mask - .append_packed_range(bit_start..bit_start + len, keep.values()); - } - } else if len != 0 { - hidden_count += set_bits.count_range(start, end); + let clear_null_rows = + |keep_mask: &mut Option, start: usize, end: usize| { keep_mask .get_or_insert_with(|| { let mut builder = BooleanBufferBuilder::new(keep.len()); - builder - .append_packed_range(keep.offset()..bit_start, keep.values()); + builder.append_packed_range( + keep.offset()..keep.offset() + start, + keep.values(), + ); builder }) - .append_n(len, false); + .append_n(end - start, false); + }; + for (row, window) in list_offsets.windows(2).enumerate() { + if list.is_valid(row) { + let start = window[0].as_usize() - first_offset; + let end = window[1].as_usize() - first_offset; + if start != covered { + hidden_count += set_bits.count_before(start) - set_before_covered; + clear_null_rows(&mut keep_mask, covered, start); + } + set_before_covered = set_bits.count_before(end); + kept_count = set_before_covered - hidden_count; + if let Some(keep_mask) = &mut keep_mask { + keep_mask.append_packed_range( + keep.offset() + start..keep.offset() + end, + keep.values(), + ); + } + covered = end; } offsets.push(O::usize_as(kept_count)); } + if covered != span_len { + clear_null_rows(&mut keep_mask, covered, span_len); + } keep_mask .map(BooleanBufferBuilder::build) .unwrap_or_else(|| keep.clone())