Add BatchCoalescer::push_filtered_batch and docs (#7652)

# Which issue does this PR close?

We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.

- Part of #7650

# Rationale for this change

In order to coalesce the result of applying a filter currently requires
first copying the results into an intermediate array (calling `filter`).

My plan is to remove this extra copy by building the final array up
directly incrementally

To do to so, there needs to be an API that can take the original data
and the filter

# What changes are included in this PR?

1. Add `BatchCoalescer::push_filtered_batch` and docs
2. Update benchmarks to use it

# Are there any user-facing changes?
New API

Author: alamb

arrow-select/src/coalesce.rs

@@ -21,13 +21,13 @@
 //! [`filter`]: crate::filter::filter
 //! [`take`]: crate::take::take
 use crate::concat::concat_batches;
-use arrow_array::StringViewArray;
+use crate::filter::filter_record_batch;
 use arrow_array::{cast::AsArray, Array, ArrayRef, RecordBatch};
+use arrow_array::{BooleanArray, StringViewArray};
 use arrow_data::ByteView;
 use arrow_schema::{ArrowError, SchemaRef};
 use std::collections::VecDeque;
 use std::sync::Arc;
-
 // Originally From DataFusion's coalesce module:
 // https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L26-L25
 
@@ -155,9 +155,62 @@ impl BatchCoalescer {
         Arc::clone(&self.schema)
     }
 
-    /// Push next batch into the Coalescer
+    /// 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`]
+    ///
+    /// # Example
+    /// # Example
+    /// ```
+    /// # use arrow_array::{record_batch, BooleanArray};
+    /// # use arrow_select::coalesce::BatchCoalescer;
+    /// let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+    /// let batch2 = record_batch!(("a", Int32, [4, 5, 6])).unwrap();
+    /// // Apply a filter to each batch to pick the first and last row
+    /// let filter = BooleanArray::from(vec![true, false, true]);
+    /// // create a new Coalescer that targets creating 1000 row batches
+    /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000);
+    /// coalescer.push_batch_with_filter(batch1, &filter);
+    /// coalescer.push_batch_with_filter(batch2, &filter);
+    /// // finsh and retrieve the created batch
+    /// coalescer.finish_buffered_batch().unwrap();
+    /// let completed_batch = coalescer.next_completed_batch().unwrap();
+    /// // filtered out 2 and 5:
+    /// let expected_batch = record_batch!(("a", Int32, [1, 3, 4, 6])).unwrap();
+    /// assert_eq!(completed_batch, expected_batch);
+    /// ```
+    pub fn push_batch_with_filter(
+        &mut self,
+        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)
+    }
+
+    /// Push all the rows from `batch` into the Coalescer
     ///
     /// See [`Self::next_completed_batch()`] to retrieve any completed batches.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow_array::record_batch;
+    /// # use arrow_select::coalesce::BatchCoalescer;
+    /// let batch1 = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
+    /// let batch2 = record_batch!(("a", Int32, [4, 5, 6])).unwrap();
+    /// // create a new Coalescer that targets creating 1000 row batches
+    /// let mut coalescer = BatchCoalescer::new(batch1.schema(), 1000);
+    /// coalescer.push_batch(batch1);
+    /// coalescer.push_batch(batch2);
+    /// // finsh and retrieve the created batch
+    /// coalescer.finish_buffered_batch().unwrap();
+    /// let completed_batch = coalescer.next_completed_batch().unwrap();
+    /// let expected_batch = record_batch!(("a", Int32, [1, 2, 3, 4, 5, 6])).unwrap();
+    /// assert_eq!(completed_batch, expected_batch);
+    /// ```
     pub fn push_batch(&mut self, batch: RecordBatch) -> Result<(), ArrowError> {
         if batch.num_rows() == 0 {
             // If the batch is empty, we don't need to do anything

arrow/benches/coalesce_kernels.rs

@@ -214,10 +214,9 @@ fn filter_streams(
     while num_output_batches > 0 {
         let filter = filter_stream.next_filter();
         let batch = data_stream.next_batch();
-        // Apply the filter to the input batch
-        let filtered_batch = arrow_select::filter::filter_record_batch(batch, filter).unwrap();
-        // Add the filtered batch to the coalescer
-        coalescer.push_batch(filtered_batch).unwrap();
+        coalescer
+            .push_batch_with_filter(batch.clone(), filter)
+            .unwrap();
         // consume (but discard) the output batch
         if coalescer.next_completed_batch().is_some() {
             num_output_batches -= 1;