comphead commented on code in PR #26507:
URL: https://github.com/apache/datafusion/pull/26507#discussion_r4318870152
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -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<O: OffsetSizeTrait>(
Review Comment:
This body is the same as `generic_list_row_has_nulls` in
`datafusion/spark/src/function/array/array_contains.rs`
(`list.values().nulls()` then `rows_have_nulls(nulls, offsets[0],
offsets[1..])`). A `GenericListArray<O>` variant next to `rows_have_nulls` in
`datafusion-common` would serve both crates and drop both wrappers.
Also, `nulls()` returns `Some` for an all-valid bitmap, in which case this
scans every word for nothing. Filtering on `nulls.null_count() != 0` (O(1))
restores the old per-row `null_count() != 0` short cut. The new benches only
use `None` for `no_nulls`, so that case is not measured.
##########
datafusion/functions-nested/src/inner_product.rs:
##########
@@ -193,15 +197,15 @@ fn general_inner_product<O: OffsetSizeTrait>(arrays:
&[ArrayRef]) -> Result<Arra
);
}
- let slice1 = values1.slice(start1, len1);
- let slice2 = values2.slice(start2, len2);
- if slice1.null_count() != 0 || slice2.null_count() != 0 {
+ let has_null =
Review Comment:
The `has_null` closure is rebuilt per row and does two bitmap lookups per
row. The same code is in `cosine_distance.rs`. OR-ing the two bitmaps once
before the loop (`Option<BooleanBuffer>` combined with `&a | &b`) gives one
lookup per row and removes the duplicated closure in both files.
##########
datafusion/functions-nested/src/inner_product.rs:
##########
@@ -212,3 +216,15 @@ fn general_inner_product<O: OffsetSizeTrait>(arrays:
&[ArrayRef]) -> Result<Arra
Ok(Arc::new(builder.finish()) as ArrayRef)
}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ #[test]
+ fn test_sliced_list() -> datafusion_common::Result<()> {
Review Comment:
Two points on the new tests.
- Sibling UDFs name this test `test_sliced_capacity` (`array_normalize`,
`array_add`, `array_scale`, `remove`). Please use the same name here and in
`cosine_distance.rs`.
- Both new tests pass the same array as both arguments, so an indexing
mistake in the second bitmap (`null_rows2`) cannot fail them.
`inner_product.slt` and `cosine_distance.slt` only cover NULL elements with
single constant rows. A multi-row `VALUES` case with a NULL element only in the
second argument, and one only in the first, would cover the per-row indexing
and belongs in the `.slt` files.
##########
datafusion/common/src/utils/mod.rs:
##########
@@ -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 {
Review Comment:
Now that this is public, two things about `new`.
1. It has no rustdoc, and the zero-padding of the final word is an invariant
the caller must uphold. `bit_chunks().iter()` without `iter_padded()` compiles,
and the counter then reads the tail word as `0` via `unwrap_or(0)`, so counts
for positions in the last partial word are silently low. A constructor taking
`&BooleanBuffer` (returning `SetBitCounter<impl Iterator<Item = u64> + '_>`)
would remove that footgun and replace the six repeated
`SetBitCounter::new(x.bit_chunks().iter_padded())` call sites.
2. The monotonic-position precondition is only a `debug_assert!(end >=
self.word_start)`, which does not catch a regression inside the current word.
By reading, in a release build a call with `end < word_start` wraps `end -
self.word_start` and the `while` loop does not terminate in practice. Please
document it on the type, or use a cheap `assert!`. I did not run this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
