Make push_batch_with_filter up to 3x faster

arrow-select/src/coalesce.rs

@@ -20,7 +20,10 @@
 //!
 //! [`filter`]: crate::filter::filter
 //! [`take`]: crate::take::take
-use crate::filter::filter_record_batch;
+use crate::filter::{
+    FilterBuilder, FilterPredicate, IndexIterator, filter_record_batch,
+    is_optimize_beneficial_record_batch,
+};
 use arrow_array::types::{BinaryViewType, StringViewType};
 use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive};
 use arrow_schema::{ArrowError, DataType, SchemaRef};
@@ -211,7 +214,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
     /// ```
@@ -237,10 +243,101 @@ 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 selective
+        if batch
+            .schema()
+            .fields()
+            .iter()
+            .any(|field| !field.data_type().is_primitive())
+            || self
+                .biggest_coalesce_batch_size
+                .map(|biggest_size| filter.true_count() > biggest_size)
+                .unwrap_or(false)
+        {
+            let batch = filter_record_batch(&batch, filter)?;
+
+            self.push_batch(batch)?;
+            return Ok(());
+        }
+
+        // Build an optimized filter predicate that chooses the best iteration strategy
+        let is_optimize_beneficial = is_optimize_beneficial_record_batch(&batch);
+        let selected_count = filter.true_count();
+
+        // 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 == batch.num_rows() {
+            return self.push_batch(batch);
+        }
+
+        let (_schema, arrays, _num_rows) = batch.into_parts();
+
+        // 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(Some(Arc::clone(array)));
+            });
+
+        // Choose iteration strategy based on the optimized predicate
+        self.copy_from_filter(filter, is_optimize_beneficial, 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: &BooleanArray,
+        is_optimize_beneficial: bool,
+        count: usize,
+    ) -> Result<(), ArrowError> {
+        let mut remaining = count;
+        let mut offset = 0;
+
+        // 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);
+
+            let sliced_filter = filter.slice(offset, to_copy);
+            let mut filter = FilterBuilder::new(&sliced_filter);
+
+            if is_optimize_beneficial {
+                filter = filter.optimize();
+            }
+            let chunk_predicate = filter.build();
+
+            // 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)?;
+            }
+
+            self.buffered_rows += to_copy;
+            offset += to_copy;
+            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 all the rows from `batch` into the Coalescer
@@ -571,6 +668,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
+        for idx in IndexIterator::new(filter.filter_array(), filter.count()) {
+            self.copy_rows(idx, 1)?;
+        }
+        Ok(())
+    }
+
     /// Finish the currently in-progress array and return it as an `ArrayRef`
     fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
 }
@@ -579,6 +685,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::{

arrow-select/src/coalesce/primitive.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::coalesce::InProgressArray;
+use crate::filter::{FilterPredicate, IndexIterator, IterationStrategy, SlicesIterator};
 use arrow_array::cast::AsArray;
 use arrow_array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
 use arrow_buffer::{NullBufferBuilder, ScalarBuffer};
@@ -92,6 +93,105 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for InProgressPrimitiveArray
         Ok(())
     }
 
+    /// Copy rows by indices 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 using slices
+                for (start, end) in SlicesIterator::new(filter.filter_array()) {
+                    self.current.extend_from_slice(&values[start..end]);
+                }
+                // Copy nulls using slices
+                if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
+                    for (start, end) in SlicesIterator::new(filter.filter_array()) {
+                        let slice = nulls.slice(start, end - start);
+                        self.nulls.append_buffer(&slice);
+                    }
+                } else {
+                    self.nulls.append_n_non_nulls(count);
+                }
+            }
+            IterationStrategy::Slices(slices) => {
+                // Copy values using precomputed slices
+                for &(start, end) in slices {
+                    self.current.extend_from_slice(&values[start..end]);
+                }
+                // Copy nulls using slices
+                if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
+                    for &(start, end) in slices {
+                        let slice = nulls.slice(start, end - start);
+                        self.nulls.append_buffer(&slice);
+                    }
+                } else {
+                    self.nulls.append_n_non_nulls(count);
+                }
+            }
+            IterationStrategy::IndexIterator => {
+                // Copy values and nulls for each index
+                if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
+                    for idx in IndexIterator::new(filter.filter_array(), count) {
+                        if nulls.is_null(idx) {
+                            self.nulls.append_null();
+                        } else {
+                            self.nulls.append_non_null();
+                        }
+                        self.current.push(values[idx]);
+                    }
+                } else {
+                    self.nulls.append_n_non_nulls(count);
+                    for idx in IndexIterator::new(filter.filter_array(), count) {
+                        self.current.push(values[idx]);
+                    }
+                }
+            }
+            IterationStrategy::Indices(indices) => {
+                // Copy values and nulls using precomputed indices
+                if let Some(nulls) = s.nulls().filter(|n| n.null_count() > 0) {
+                    for &idx in indices {
+                        if nulls.is_null(idx) {
+                            self.nulls.append_null();
+                        } else {
+                            self.nulls.append_non_null();
+                        }
+                    }
+                } else {
+                    self.nulls.append_n_non_nulls(count);
+                }
+                self.current.extend(indices.iter().map(|&idx| values[idx]));
+            }
+            IterationStrategy::All => {
+                // Copy all values
+                self.current.extend_from_slice(values);
+                if let Some(nulls) = s.nulls() {
+                    self.nulls.append_buffer(nulls);
+                } else {
+                    self.nulls.append_n_non_nulls(values.len());
+                }
+            }
+            IterationStrategy::None => {
+                // Nothing to copy
+            }
+        }
+
+        Ok(())
+    }
+
     fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
         // take and reset the current values and nulls
         let values = std::mem::take(&mut self.current);

arrow-select/src/filter.rs

@@ -80,13 +80,17 @@ impl Iterator for SlicesIterator<'_> {
 ///
 /// This provides the best performance on most predicates, apart from those which keep
 /// large runs and therefore favour [`SlicesIterator`]
-struct IndexIterator<'a> {
+pub struct IndexIterator<'a> {
     remaining: usize,
     iter: BitIndexIterator<'a>,
 }
 
 impl<'a> IndexIterator<'a> {
-    fn new(filter: &'a BooleanArray, remaining: usize) -> Self {
+    /// Creates a new [`IndexIterator`] from a [`BooleanArray`]
+    ///
+    /// # Panics
+    /// Panics if `filter` has null values
+    pub fn new(filter: &'a BooleanArray, remaining: usize) -> Self {
         assert_eq!(filter.null_count(), 0);
         let iter = filter.values().set_indices();
         Self { remaining, iter }
@@ -167,6 +171,21 @@ pub fn filter(values: &dyn Array, predicate: &BooleanArray) -> Result<ArrayRef,
     filter_array(values, &predicate)
 }
 
+/// Determines if calling [FilterBuilder::optimize] is beneficial for the
+/// given [`RecordBatch`].
+pub fn is_optimize_beneficial_record_batch(record_batch: &RecordBatch) -> bool {
+    let num_cols = record_batch.num_columns();
+    if num_cols > 1 {
+        return true;
+    }
+    if num_cols == 1 {
+        return FilterBuilder::is_optimize_beneficial(
+            record_batch.schema_ref().field(0).data_type(),
+        );
+    }
+    false
+}
+
 /// Returns a filtered [RecordBatch] where the corresponding elements of
 /// `predicate` are true.
 ///
@@ -182,13 +201,7 @@ pub fn filter_record_batch(
     predicate: &BooleanArray,
 ) -> Result<RecordBatch, ArrowError> {
     let mut filter_builder = FilterBuilder::new(predicate);
-    let num_cols = record_batch.num_columns();
-    if num_cols > 1
-        || (num_cols > 0
-            && FilterBuilder::is_optimize_beneficial(
-                record_batch.schema_ref().field(0).data_type(),
-            ))
-    {
+    if is_optimize_beneficial_record_batch(record_batch) {
         // Only optimize if filtering more than one column or if the column contains multiple internal arrays
         // Otherwise, the overhead of optimization can be more than the benefit
         filter_builder = filter_builder.optimize();
@@ -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 high selectivity filters (> 80% of rows selected)
     SlicesIterator,
     /// A lazily evaluated iterator of indices
+    ///
+    /// Best for low selectivity filters (< 80% of rows selected)
     IndexIterator,
     /// A precomputed list of indices
     Indices(Vec<usize>),
-    /// A precomputed array of ranges
+    /// A precomputed array of ranges (start, end)
     Slices(Vec<(usize, usize)>),
     /// Select all rows
     All,
@@ -295,7 +315,13 @@ enum IterationStrategy {
 impl IterationStrategy {
     /// The default [`IterationStrategy`] for a filter of length `filter_length`
     /// and selecting `filter_count` rows
-    fn default_strategy(filter_length: usize, filter_count: usize) -> Self {
+    ///
+    /// Returns:
+    /// - [`IterationStrategy::None`] if `filter_count` is 0
+    /// - [`IterationStrategy::All`] if `filter_count == filter_length`
+    /// - [`IterationStrategy::SlicesIterator`] if selectivity > 80%
+    /// - [`IterationStrategy::IndexIterator`] otherwise
+    pub fn default_strategy(filter_length: usize, filter_count: usize) -> Self {
         if filter_length == 0 || filter_count == 0 {
             return IterationStrategy::None;
         }
@@ -359,6 +385,16 @@ impl FilterPredicate {
     pub fn count(&self) -> usize {
         self.count
     }
+
+    /// Returns the iteration strategy used by this [`FilterPredicate`]
+    pub fn strategy(&self) -> &IterationStrategy {
+        &self.strategy
+    }
+
+    /// Returns the underlying filter array
+    pub fn filter_array(&self) -> &BooleanArray {
+        &self.filter
+    }
 }
 
 fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result<ArrayRef, ArrowError> {