alamb commented on code in PR #8951:
URL: https://github.com/apache/arrow-rs/pull/8951#discussion_r2615576833
##########
arrow-select/src/coalesce.rs:
##########
@@ -526,6 +631,118 @@ impl BatchCoalescer {
}
}
+/// Find the position after the n-th set bit in a boolean array starting from
`start`.
+/// Returns the position after the n-th set bit, or the end of the array if
fewer than n bits are set.
Review Comment:
I recommend we move this code into BooleanBuffer as well so it is easier to
find / reuse
##########
arrow-buffer/src/builder/null.rs:
##########
@@ -402,4 +481,43 @@ mod tests {
assert_eq!(builder.finish(), None);
}
+
+ #[test]
+ fn test_extend() {
+ // Test small extend (less than 64 bits)
+ let mut builder = NullBufferBuilder::new(0);
+ builder.extend([true, false, true, true].iter().copied());
+ // bits: 0=true, 1=false, 2=true, 3=true -> 0b1101 = 13
+ assert_eq!(builder.as_slice().unwrap(), &[0b1101_u8]);
+
+ // Test extend with exactly 64 bits
+ let mut builder = NullBufferBuilder::new(0);
+ let pattern: Vec<bool> = (0..64).map(|i| i % 2 == 0).collect();
+ builder.extend(pattern.iter().copied());
+ // Even positions are true: 0, 2, 4, ... -> bits 0, 2, 4, ...
+ // In little-endian: 0b01010101 repeated
+ assert_eq!(
+ builder.as_slice().unwrap(),
+ &[0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55]
+ );
+
+ // Test extend with more than 64 bits (tests chunking)
+ let mut builder = NullBufferBuilder::new(0);
+ let pattern: Vec<bool> = (0..100).map(|i| i % 3 == 0).collect();
+ builder.extend(pattern.iter().copied());
+ assert_eq!(builder.len(), 100);
+ // Verify a few specific bits
+ let buf = builder.finish().unwrap();
+ assert!(buf.is_valid(0)); // 0 % 3 == 0
+ assert!(!buf.is_valid(1)); // 1 % 3 != 0
+ assert!(!buf.is_valid(2)); // 2 % 3 != 0
+ assert!(buf.is_valid(3)); // 3 % 3 == 0
+ assert!(buf.is_valid(99)); // 99 % 3 == 0
+
+ // Test extend with non-aligned start (tests bit-by-bit path)
+ let mut builder = NullBufferBuilder::new(0);
+ builder.append_non_null(); // Start at bit 1 (non-aligned)
+ builder.extend([false, true, false, true].iter().copied());
Review Comment:
I think we should probably test non aligned writes with more than 64 bits as
well (this only copies 4 bits)
##########
arrow-select/src/coalesce/primitive.rs:
##########
@@ -92,6 +93,145 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
Ok(())
}
+ /// Copy rows using a predicate
+ fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(),
ArrowError> {
+ self.ensure_capacity();
+
+ let s = self
+ .source
+ .as_ref()
+ .ok_or_else(|| {
+ ArrowError::InvalidArgumentError(
+ "Internal Error: InProgressPrimitiveArray: source not
set".to_string(),
+ )
+ })?
+ .as_primitive::<T>();
+
+ let values = s.values();
+ let count = filter.count();
+
+ // Use the predicate's strategy for optimal iteration
+ match filter.strategy() {
+ IterationStrategy::SlicesIterator => {
+ // Copy values, nulls using slices
+ if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
+ for (start, end) in
SlicesIterator::new(filter.filter_array()) {
+ // SAFETY: slices are derived from filter predicate
+ self.current
+ .extend_from_slice(unsafe {
values.get_unchecked(start..end) });
+ let slice = nulls.slice(start, end - start);
+ self.nulls.append_buffer(&slice);
+ }
+ } else {
+ for (start, end) in
SlicesIterator::new(filter.filter_array()) {
+ // SAFETY: SlicesIterator produces valid ranges
derived from filter
+ self.current
+ .extend_from_slice(unsafe {
values.get_unchecked(start..end) });
+ }
+ self.nulls.append_n_non_nulls(count);
+ }
+ }
+ IterationStrategy::Slices(slices) => {
Review Comment:
I think this function needs some tests
I ran code coverage like this
```shell
cargo llvm-cov test --html -p arrow-buffer -p arrow-select
```
And there appears to be no coverage
<img width="1097" height="949" alt="Image"
src="https://github.com/user-attachments/assets/0e644657-d723-4f5f-925a-3c4ce86a78bc";
/>
##########
arrow-buffer/src/builder/null.rs:
##########
@@ -193,6 +193,85 @@ impl NullBufferBuilder {
}
}
+ /// Extends this builder with validity values.
+ ///
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::NullBufferBuilder;
+ /// let mut builder = NullBufferBuilder::new(8);
+ /// let validities = [true, false, true, true];
+ /// builder.extend(validities.iter().copied());
+ /// assert_eq!(builder.len(), 4);
+ /// ```
+ pub fn extend<I: Iterator<Item = bool>>(&mut self, iter: I) {
+ let (lower, upper) = iter.size_hint();
+ let len = upper.expect("Iterator must have exact size_hint");
+ debug_assert_eq!(lower, len, "Iterator must have exact size_hint");
+
+ if len == 0 {
+ return;
+ }
+
+ // Materialize since we're about to append bits
+ self.materialize_if_needed();
+
+ let buf = self.bitmap_builder.as_mut().unwrap();
+ let start_len = buf.len();
+ // Advance to allocate space, initializing new bits to 0
+ buf.advance(len);
+
+ let slice = buf.as_slice_mut();
+ let mut bit_idx = start_len;
+ let end_bit = start_len + len;
+
+ // Process in chunks of 64 bits when byte-aligned for better
performance
+ if start_len % 8 == 0 {
+ let start_byte = start_len / 8;
+ let mut iter = iter.peekable();
+
+ // Process full u64 chunks (64 bits at a time)
+ while bit_idx + 64 <= end_bit && iter.peek().is_some() {
+ let mut chunk: u64 = 0;
+ for i in 0..64 {
+ if let Some(valid) = iter.next() {
+ if valid {
+ chunk |= 1u64 << i;
+ }
+ } else {
+ break;
+ }
+ }
+ let byte_idx = (bit_idx - start_len) / 8 + start_byte;
+ // Write the u64 chunk as 8 bytes
+ slice[byte_idx..byte_idx +
8].copy_from_slice(&chunk.to_le_bytes());
Review Comment:
could try unsafe here too as you ensured the right length above
##########
arrow-buffer/src/builder/null.rs:
##########
@@ -193,6 +193,85 @@ impl NullBufferBuilder {
}
}
+ /// Extends this builder with validity values.
+ ///
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::NullBufferBuilder;
+ /// let mut builder = NullBufferBuilder::new(8);
+ /// let validities = [true, false, true, true];
+ /// builder.extend(validities.iter().copied());
+ /// assert_eq!(builder.len(), 4);
+ /// ```
+ pub fn extend<I: Iterator<Item = bool>>(&mut self, iter: I) {
+ let (lower, upper) = iter.size_hint();
Review Comment:
I think this method would be more generally when appending to any
BooleanBuffer rather than just NullBufferBuilder
As part of the goal to consolidate mutable boolean operations in
`BooleanBufferBuilder` so it is easier to find (and optimize) them, would you
be willing to move this code to `BooleanBufferBuilder` so that the code in
NullBufferBuilder looks like something like this (which is what most other
methods in NullBufferBuilder look like)?
```rust
pub fn extend<I: Iterator<Item = bool>>(&mut self, iter: I) {
// Materialize since we're about to append bits
self.materialize_if_needed();
self.bitmap_builder.as_mut().unwrap().extend(iter)
}
```
##########
arrow-buffer/src/builder/null.rs:
##########
@@ -193,6 +193,85 @@ impl NullBufferBuilder {
}
}
+ /// Extends this builder with validity values.
+ ///
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::NullBufferBuilder;
+ /// let mut builder = NullBufferBuilder::new(8);
+ /// let validities = [true, false, true, true];
+ /// builder.extend(validities.iter().copied());
+ /// assert_eq!(builder.len(), 4);
+ /// ```
+ pub fn extend<I: Iterator<Item = bool>>(&mut self, iter: I) {
+ let (lower, upper) = iter.size_hint();
+ let len = upper.expect("Iterator must have exact size_hint");
+ debug_assert_eq!(lower, len, "Iterator must have exact size_hint");
+
+ if len == 0 {
+ return;
+ }
+
+ // Materialize since we're about to append bits
+ self.materialize_if_needed();
+
+ let buf = self.bitmap_builder.as_mut().unwrap();
+ let start_len = buf.len();
+ // Advance to allocate space, initializing new bits to 0
+ buf.advance(len);
+
+ let slice = buf.as_slice_mut();
+ let mut bit_idx = start_len;
+ let end_bit = start_len + len;
+
+ // Process in chunks of 64 bits when byte-aligned for better
performance
+ if start_len % 8 == 0 {
+ let start_byte = start_len / 8;
+ let mut iter = iter.peekable();
+
+ // Process full u64 chunks (64 bits at a time)
+ while bit_idx + 64 <= end_bit && iter.peek().is_some() {
Review Comment:
As a follow on PR, it might be worth aligning first on 64 bit boundaries (so
the underlying code doesn't have to handle aligning) -- aka handle bits 0..63
(until 64 bit alignment) specially and then use the u64 path
##########
arrow-select/src/coalesce.rs:
##########
@@ -237,10 +243,109 @@ impl BatchCoalescer {
batch: RecordBatch,
filter: &BooleanArray,
) -> Result<(), ArrowError> {
- // TODO: optimize this to avoid materializing (copying the results
Review Comment:
🎉
##########
arrow-select/src/filter.rs:
##########
@@ -276,15 +289,22 @@ impl FilterBuilder {
}
/// The iteration strategy used to evaluate [`FilterPredicate`]
-#[derive(Debug)]
-enum IterationStrategy {
- /// A lazily evaluated iterator of ranges
+///
+/// This determines how the filter will iterate over the selected rows.
+/// The strategy is chosen based on the selectivity of the filter.
+#[derive(Debug, Clone)]
+pub enum IterationStrategy {
+ /// A lazily evaluated iterator of ranges (slices)
+ ///
+ /// Best for low selectivity filters
Review Comment:
I think the term "low selectivity" filters is inconsistently used. I
recommend making this explicit
```suggestion
/// Best for low selectivity filters (which select a relatively large
number of rows)
```
##########
arrow-select/src/filter.rs:
##########
@@ -276,15 +289,22 @@ impl FilterBuilder {
}
/// The iteration strategy used to evaluate [`FilterPredicate`]
-#[derive(Debug)]
-enum IterationStrategy {
- /// A lazily evaluated iterator of ranges
+///
+/// This determines how the filter will iterate over the selected rows.
+/// The strategy is chosen based on the selectivity of the filter.
+#[derive(Debug, Clone)]
+pub enum IterationStrategy {
+ /// A lazily evaluated iterator of ranges (slices)
+ ///
+ /// Best for low selectivity filters
SlicesIterator,
/// A lazily evaluated iterator of indices
+ ///
+ /// Best for high selectivity filters
Review Comment:
```suggestion
/// Best for high selectivity filters (which select a relatively low
number of rows)
```
##########
arrow-select/src/coalesce.rs:
##########
@@ -571,6 +788,15 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
/// Return an error if the source array is not set
fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(),
ArrowError>;
+ /// Copy rows at the given indices from the current source array into the
in-progress array
+ fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(),
ArrowError> {
+ // Default implementation: iterate over indices from the filter
Review Comment:
It seems like as a follow on we should implement something similar for the
byte array filter types? If that is true I'll file a ticket
--
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]
