Make push_batch_with_filter up to 3x faster

arrow-buffer/src/builder/null.rs

@@ -196,6 +196,33 @@ impl NullBufferBuilder {
         }
     }
 
+    /// Extends this builder with validity values.
+    ///
+    /// # Safety
+    /// The caller must ensure that the iterator reports the correct length.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow_buffer::NullBufferBuilder;
+    /// let mut builder = NullBufferBuilder::new(8);
+    /// let validities = [true, false, true, true];
+    /// unsafe { builder.extend_trusted_len(validities.iter().copied()); }
+    /// assert_eq!(builder.len(), 4);
+    /// ```
+    pub unsafe fn extend_trusted_len<I: Iterator<Item = bool>>(&mut self, iter: I) {
+        // Materialize since we're about to append bits
+        self.materialize_if_needed();
+
+        unsafe {
+            self.bitmap_builder
+                .as_mut()
+                .unwrap()
+                .extend_trusted_len(iter)
+        };
+    }
+
+    /// Builds the null buffer and resets the builder.
+    /// Returns `None` if the builder only contains `true`s.
     /// Builds the [`NullBuffer`] and resets the builder.
     ///
     /// Returns `None` if the builder only contains `true`s. Use [`Self::build`]
@@ -412,4 +439,47 @@ mod tests {
 
         assert_eq!(builder.finish(), None);
     }
+
+    #[test]
+    fn test_extend() {
+        // Test small extend (less than 64 bits)
+        let mut builder = NullBufferBuilder::new(0);
+        unsafe {
+            builder.extend_trusted_len([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();
+        unsafe {
+            builder.extend_trusted_len(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();
+        unsafe { builder.extend_trusted_len(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)
+        unsafe { builder.extend_trusted_len([false, true, false, true].iter().copied()) };
+        assert_eq!(builder.as_slice().unwrap(), &[0b10101_u8]);
+    }
 }

arrow-select/src/coalesce.rs

@@ -20,7 +20,8 @@
 //!
 //! [`filter`]: crate::filter::filter
 //! [`take`]: crate::take::take
-use crate::filter::filter_record_batch;
+use crate::filter::{FilterBuilder, FilterPredicate, is_optimize_beneficial_record_batch};
+
 use crate::take::take_record_batch;
 use arrow_array::types::{BinaryViewType, StringViewType};
 use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive};
@@ -212,7 +213,10 @@ impl BatchCoalescer {
     /// Push a batch into the Coalescer after applying a filter
     ///
     /// This is semantically equivalent of calling [`Self::push_batch`]
-    /// with the results from  [`filter_record_batch`]
+    /// with the results from  [`filter_record_batch`], but avoids
+    /// materializing the intermediate filtered batch.
+    ///
+    /// [`filter_record_batch`]: crate::filter::filter_record_batch
     ///
     /// # Example
     /// ```
@@ -238,10 +242,103 @@ impl BatchCoalescer {
         batch: RecordBatch,
         filter: &BooleanArray,
     ) -> Result<(), ArrowError> {
-        // TODO: optimize this to avoid materializing (copying the results
-        // of filter to a new batch)
-        let filtered_batch = filter_record_batch(&batch, filter)?;
-        self.push_batch(filtered_batch)
+        // We only support primitve now, fallback to filter_record_batch for other types
+        // Also, skip optimization when filter is not very selectivex§
+
+        // Build an optimized filter predicate that chooses the best iteration strategy
+        // Byteview does use a filter as part of calculating ideal buffer sizes, so optimizing is helpful even for
+        // a single array
+        let is_optimize_beneficial = is_optimize_beneficial_record_batch(&batch)
+            || batch.columns().len() == 1
+                && matches!(
+                    batch.columns()[0].data_type(),
+                    DataType::BinaryView | DataType::Utf8View
+                );
+        let selected_count = filter.true_count();
+        let num_rows = batch.num_rows();
+
+        // Fast path: skip if no rows selected
+        if selected_count == 0 {
+            return Ok(());
+        }
+
+        // Fast path: if all rows selected, just push the batch
+        if selected_count == num_rows {
+            return self.push_batch(batch);
+        }
+
+        let (_schema, arrays, _num_rows) = batch.into_parts();
+
+        let mut filter_builder = FilterBuilder::new(&filter);
+
+        if is_optimize_beneficial {
+            filter_builder = filter_builder.optimize();
+        }
+
+        let filter = filter_builder.build();
+        // Setup input arrays as sources
+        assert_eq!(arrays.len(), self.in_progress_arrays.len());
+        self.in_progress_arrays
+            .iter_mut()
+            .zip(arrays)
+            .for_each(|(in_progress, array)| {
+                in_progress.set_source_from_filter(Some(array), &filter);
+            });
+
+        // Choose iteration strategy based on the optimized predicate
+        self.copy_from_filter(filter, selected_count)?;
+        // Clear sources to allow memory to be freed
+        for in_progress in self.in_progress_arrays.iter_mut() {
+            in_progress.set_source(None);
+        }
+
+        Ok(())
+    }
+
+    /// Helper to copy rows at the given indices, handling batch boundaries efficiently
+    ///
+    /// This method batches the index iteration to avoid per-row batch boundary checks.
+    fn copy_from_filter(
+        &mut self,
+        filter: FilterPredicate,
+        count: usize,
+    ) -> Result<(), ArrowError> {
+        let mut remaining = count;
+        let mut filter_pos = 0; // Position in the filter array
+
+        // Build an optimized filter predicate once for the whole input batch
+
+        // We need to process the filter in chunks that fit the target batch size
+        while remaining > 0 {
+            let space_in_batch = self.target_batch_size - self.buffered_rows;
+            let to_copy = remaining.min(space_in_batch);
+
+            // Find how many filter positions we need to cover `to_copy` set bits
+            // Skip the expensive search if all remaining rows fit in the current batch
+            let chunk_len = if remaining <= space_in_batch {
+                filter.len() - filter_pos
+            } else {
+                filter.find_nth_set_bit_position(filter_pos, to_copy) - filter_pos
+            };
+
+            let chunk_predicate = filter.slice_with_count(filter_pos, chunk_len, to_copy);
+
+            // Copy all collected indices in one call per array
+            for in_progress in self.in_progress_arrays.iter_mut() {
+                in_progress.copy_rows_by_filter(&chunk_predicate, filter_pos, chunk_len)?;
+            }
+
+            self.buffered_rows += to_copy;
+            filter_pos += chunk_len;
+            remaining -= to_copy;
+
+            // If we've filled the batch, finish it
+            if self.buffered_rows >= self.target_batch_size {
+                self.finish_buffered_batch()?;
+            }
+        }
+
+        Ok(())
     }
 
     /// Push a batch into the Coalescer after applying a set of indices
@@ -598,13 +695,31 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
     /// current in-progress array
     fn set_source(&mut self, source: Option<ArrayRef>);
 
+    /// Set the source array with a filter, allowing for calculating GC based on filter
+    ///
+    /// Default implementation just calls [`Self::set_source`]
+    fn set_source_from_filter(&mut self, source: Option<ArrayRef>, _filter: &FilterPredicate) {
+        self.set_source(source);
+    }
+
     /// Copy rows from the current source array into the in-progress array
     ///
     /// The source array is set by [`Self::set_source`].
     ///
     /// Return an error if the source array is not set
     fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>;
 
+    /// Copy rows from the source array between the specified offset and len that
+    /// match the predicate to the output array
+    ///
+    /// TODO add an example
+    fn copy_rows_by_filter(
+        &mut self,
+        filter: &FilterPredicate,
+        offset: usize,
+        len: usize,
+    ) -> Result<(), ArrowError>;
+
     /// Finish the currently in-progress array and return it as an `ArrayRef`
     fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
 }
@@ -613,6 +728,7 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
 mod tests {
     use super::*;
     use crate::concat::concat_batches;
+    use crate::filter::filter_record_batch;
     use arrow_array::builder::StringViewBuilder;
     use arrow_array::cast::AsArray;
     use arrow_array::types::Int32Type;