perf: Optimize repeated counting of set bits

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<O: ArrowNativeType>(offsets: &OffsetBuffer<O>) -> (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<I> {
+    /// 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<I: Iterator<Item = u64>> SetBitCounter<I> {
+    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<Item = usize>,
+) -> 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::<i32>::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::<Vec<_>>();
+        assert_eq!(has_nulls.iter().collect::<Vec<_>>(), 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 =

datafusion/functions-nested/Cargo.toml

@@ -157,3 +157,7 @@ name = "cardinality"
 [[bench]]
 harness = false
 name = "range"
+
+[[bench]]
+harness = false
+name = "vector_math"

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<ArrayRef>) {
+    let arg_types = inputs
+        .iter()
+        .map(|input| input.data_type().clone())
+        .collect::<Vec<_>>();
+    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::<Vec<f64>>();
+    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);

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<O: OffsetSizeTrait>(
     list_array: &GenericListArray<O>,
     values_nulls: &NullBuffer,
 ) -> Result<(OffsetBuffer<O>, 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::<i32>(), &keep_mask, kept_count)
@@ -264,7 +235,7 @@ fn compact_non_primitive<O: OffsetSizeTrait>(
         }
         _ => 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