This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new a19fc628b9 Add `BatchCoalescer::push_filtered_batch` and docs (#7652)
a19fc628b9 is described below
commit a19fc628b96371b194c0e9c89174c3209adcd195
Author: Andrew Lamb <[email protected]>
AuthorDate: Tue Jun 17 07:37:39 2025 -0400
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 https://github.com/apache/arrow-rs/pull/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
---
arrow-select/src/coalesce.rs | 59 +++++++++++++++++++++++++++++++++++++--
arrow/benches/coalesce_kernels.rs | 7 ++---
2 files changed, 59 insertions(+), 7 deletions(-)
diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs
index c807d99d77..fb126aa9e6 100644
--- a/arrow-select/src/coalesce.rs
+++ b/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
diff --git a/arrow/benches/coalesce_kernels.rs
b/arrow/benches/coalesce_kernels.rs
index 16db07e388..1168d4b023 100644
--- a/arrow/benches/coalesce_kernels.rs
+++ b/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;
