jayzhan211 commented on code in PR #25503:
URL: https://github.com/apache/datafusion/pull/25503#discussion_r4056922244
##########
datafusion/physical-plan/src/coalesce/mod.rs:
##########
@@ -120,6 +120,44 @@ impl LimitedBatchCoalescer {
Ok(PushBatchStatus::Continue)
}
+ /// Pushes the next [`RecordBatch`] into the coalescer after applying
`filter`,
+ /// avoiding a separate materialization pass compared to calling
+ /// [`filter_record_batch`] followed by [`Self::push_batch`].
+ ///
+ /// [`filter_record_batch`]: arrow::compute::filter_record_batch
+ pub fn push_batch_with_filter(
+ &mut self,
+ batch: RecordBatch,
+ filter: &BooleanArray,
+ ) -> Result<PushBatchStatus> {
+ assert_or_internal_err!(
+ !self.finished,
+ "LimitedBatchCoalescer: cannot push batch after finish"
+ );
+
+ let Some(fetch) = self.fetch else {
+ self.inner.push_batch_with_filter(batch, filter)?;
+ return Ok(PushBatchStatus::Continue);
+ };
+
+ if self.total_rows >= fetch {
+ return Ok(PushBatchStatus::LimitReached);
+ }
+
+ let selected_count = filter.true_count();
+ if self.total_rows + selected_count >= fetch {
+ let remaining = fetch - self.total_rows;
+ let truncated = first_n_true(filter, remaining);
+ self.total_rows += remaining;
+ self.inner.push_batch_with_filter(batch, &truncated)?;
Review Comment:
```suggestion
let mask = match filter.null_count() {
0 => filter.clone(),
_ => prep_null_mask_filter(filter),
};
// one past the row holding the `remaining`-th selected value
let end = mask
.values()
.set_indices()
.nth(remaining - 1)
.map_or(0, |i| i + 1);
self.total_rows += remaining;
self.inner
.push_batch_with_filter(batch.slice(0, end), &mask.slice(0,
end))?;
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
