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/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); 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/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 73bf3cbf4c84f..5201ddeb0e3a8 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::{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; @@ -79,35 +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 { - let set_before_row = set_bits.count_before(start); - hidden_count += set_bits.count_before(end) - set_before_row; + 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()) @@ -120,45 +139,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. @@ -631,6 +611,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::*; @@ -758,21 +752,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() { @@ -874,7 +853,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) 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); 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),