mbutrovich commented on code in PR #26507:
URL: https://github.com/apache/datafusion/pull/26507#discussion_r4340403894
##########
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> {
Review Comment:
Should this move to arrow-rs eventually? As far as I can tell, arrow-rs has
no counterpart. Its closest APIs answer one range per call.
[`Buffer::count_set_bits_offset`](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/immutable.rs#L370-L374)
builds a new `UnalignedBitChunk` on every call, which is the overhead this
type avoids.
[`BooleanBuffer::find_nth_set_bit_position`](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/boolean.rs#L440-L452)
answers the inverse question (where the n-th set bit is) and also rescans from
`start` on every call. A counter like this one, or a kernel that takes an
`OffsetBuffer` and returns per-row counts, would fit next to them in
`arrow-buffer`. arrow-rs can ship an additive API like that in a [minor
release](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/README.md?plain=1#L80-L82).
I'm happy to land it here first. Could you open an arrow-rs issue for the
upstream API, and a DataFusion issue to switch to it once it lands? Because
`SetBitCounter` becomes public in `datafusion-common`, switching later means a
deprecation cycle. That makes the constructor question in @comphead's thread
worth settling with the upstream API in mind. arrow-rs exposes its bitmap
iterators as `&self` methods on `BooleanBuffer`
([`bit_chunks`](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/boolean.rs#L458),
[`set_indices`](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/boolean.rs#L622),
[`set_slices`](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/boolean.rs#L632)),
so a constructor that takes `&BooleanBuffer` would also match how this would
look upstream.
##########
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:
> The new benches only use `None` for `no_nulls`, so that case is not
measured.
With `None`, `rows_with_null_elements` returns before calling
`rows_have_nulls`, so the `no_nulls` cases in `vector_math` never run
`SetBitCounter`. As I read the diff, their speedup comes from replacing the
per-row `values.slice(start, len)` with `&values.values()[start..end]`. Could
`float_lists` also build an all-valid `NullBuffer` as a third case? That would
measure the path this comment is about, and the counter's cost when nothing is
NULL.
--
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]
