diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..70c86368c1477 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; -use arrow::compute::BatchCoalescer; +use arrow::array::{Array, BooleanArray, RecordBatch}; +use arrow::compute::{BatchCoalescer, prep_null_mask_filter}; use arrow::datatypes::SchemaRef; use datafusion_common::{Result, assert_or_internal_err}; @@ -120,6 +120,53 @@ 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 { + 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 mask = match filter.null_count() { + 0 => filter.clone(), + _ => prep_null_mask_filter(filter), + }; + 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))?; + return Ok(PushBatchStatus::LimitReached); + } + + self.total_rows += selected_count; + self.inner.push_batch_with_filter(batch, filter)?; + Ok(PushBatchStatus::Continue) + } + /// Return true if there is no data buffered pub fn is_empty(&self) -> bool { self.inner.is_empty() @@ -225,6 +272,94 @@ mod tests { .run() } + #[test] + fn test_push_batch_with_filter_nulls_and_fetch() { + let batch = uint32_batch(0..8); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(3)); + let filter = BooleanArray::from(vec![ + None, + Some(true), + None, + Some(false), + Some(true), + Some(true), + None, + Some(true), + ]); + + assert_eq!( + coalescer.push_batch_with_filter(batch, &filter).unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![1, 4, 5]); + } + + #[test] + fn test_push_batch_with_filter_fetch_boundaries() { + let batch1 = uint32_batch(0..4); + let batch2 = uint32_batch(4..8); + let mut coalescer = LimitedBatchCoalescer::new(batch1.schema(), 100, Some(3)); + + assert_eq!( + coalescer + .push_batch_with_filter( + batch1, + &BooleanArray::from(vec![true, false, true, false]), + ) + .unwrap(), + PushBatchStatus::Continue, + ); + assert_eq!( + coalescer + .push_batch_with_filter( + batch2, + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![0, 2, 4]); + + let batch = uint32_batch(0..4); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(2)); + assert_eq!( + coalescer + .push_batch_with_filter( + batch, + &BooleanArray::from(vec![true, false, true, false]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + assert_eq!( + coalescer + .push_batch_with_filter( + uint32_batch(4..8), + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert_next_batch_values(&mut coalescer, vec![0, 2]); + + let batch = uint32_batch(0..4); + let mut coalescer = LimitedBatchCoalescer::new(batch.schema(), 100, Some(0)); + assert_eq!( + coalescer + .push_batch_with_filter( + batch, + &BooleanArray::from(vec![true, true, true, true]), + ) + .unwrap(), + PushBatchStatus::LimitReached, + ); + coalescer.finish().unwrap(); + assert!(coalescer.next_completed_batch().is_none()); + } + /// Test for [`LimitedBatchCoalescer`] /// /// Pushes the input batches to the coalescer and verifies that the resulting @@ -367,6 +502,15 @@ mod tests { .unwrap() } + fn assert_next_batch_values( + coalescer: &mut LimitedBatchCoalescer, + expected: Vec, + ) { + let output = coalescer.next_completed_batch().unwrap(); + let expected = UInt32Array::from(expected); + assert_eq!(output.column(0).as_ref(), &expected as &dyn Array); + } + fn batch_to_pretty_strings(batch: &RecordBatch) -> String { arrow::util::pretty::pretty_format_batches(std::slice::from_ref(batch)) .unwrap() diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 7377ec893621c..f71d67d04e39c 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1466,9 +1466,7 @@ impl Stream for FilterExecStream { match as_boolean_array(&array) { Ok(filter_array) => { self.metrics.selectivity.add_total(batch.num_rows()); - // TODO: support push_batch_with_filter in LimitedBatchCoalescer - let batch = filter_record_batch(&batch, filter_array)?; - let state = self.batch_coalescer.push_batch(batch)?; + let state = self.batch_coalescer.push_batch_with_filter(batch, filter_array)?; Ok(state) } Err(_) => { @@ -1571,13 +1569,51 @@ pub type EqualAndNonEqual<'a> = #[cfg(test)] mod tests { use super::*; + use crate::common::collect; use crate::empty::EmptyExec; use crate::expressions::*; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::Int32Array; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + #[tokio::test] + async fn test_filter_exec_fetch_truncates_within_selected_rows() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + )?; + let input = test::TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&schema), + None, + )?; + let predicate = binary(col("i", &schema)?, Operator::GtEq, lit(2i32), &schema)?; + let filter = Arc::new( + FilterExecBuilder::new(predicate, input) + .with_fetch(Some(3)) + .build()?, + ); + + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(filter.execute(0, task_ctx)?).await?; + let values: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, vec![2, 3, 4]); + Ok(()) + } + #[test] fn filter_rejects_zero_batch_size() -> Result<()> { let input: Arc = @@ -2348,7 +2384,7 @@ mod tests { // A mathematical lower bound of 5 would exclude a valid input value. let batch = RecordBatch::try_new( input.schema(), - vec![Arc::new(arrow::array::Int32Array::from(vec![i32::MIN]))], + vec![Arc::new(Int32Array::from(vec![i32::MIN]))], ) .unwrap(); let result = predicate.evaluate(&batch).unwrap().into_array(1).unwrap();