diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..54873dc124f8d 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -125,12 +125,26 @@ impl LimitedBatchCoalescer { self.inner.is_empty() } + /// Complete the current buffered batch without finishing the coalescer. + /// + /// Subsequent calls to [`Self::push_batch`] are allowed. + pub(crate) fn flush(&mut self) -> Result<()> { + assert_or_internal_err!( + !self.finished, + "LimitedBatchCoalescer: cannot flush after finish" + ); + self.inner.finish_buffered_batch()?; + Ok(()) + } + /// Complete the current buffered batch and finish the coalescer /// /// Any subsequent calls to `push_batch()` will return an Err pub fn finish(&mut self) -> Result<()> { - self.inner.finish_buffered_batch()?; - self.finished = true; + if !self.finished { + self.flush()?; + self.finished = true; + } Ok(()) } @@ -225,6 +239,37 @@ mod tests { .run() } + #[test] + fn test_flush_allows_subsequent_batches() { + let schema = uint32_batch(0..1).schema(); + let mut coalescer = LimitedBatchCoalescer::new(schema, 10, None); + + assert_eq!( + coalescer.push_batch(uint32_batch(0..3)).unwrap(), + PushBatchStatus::Continue + ); + assert!(coalescer.next_completed_batch().is_none()); + coalescer.flush().unwrap(); + assert_eq!( + coalescer.next_completed_batch().unwrap(), + uint32_batch(0..3) + ); + + assert_eq!( + coalescer.push_batch(uint32_batch(3..5)).unwrap(), + PushBatchStatus::Continue + ); + coalescer.flush().unwrap(); + assert_eq!( + coalescer.next_completed_batch().unwrap(), + uint32_batch(3..5) + ); + + coalescer.finish().unwrap(); + 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 diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 12771eec78470..d6cd0356a8cef 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -621,19 +621,35 @@ impl ExecutionPlan for FilterExec { context.session_id(), context.task_id() ); - let metrics = FilterExecMetrics::new(&self.metrics, partition); - Ok(Box::pin(FilterExecStream { - schema: self.schema(), - predicate: Arc::clone(&self.predicate), - input: self.input.execute(partition, context)?, - metrics, - projection: self.projection.clone(), - batch_coalescer: LimitedBatchCoalescer::new( - self.schema(), - self.batch_size, - self.fetch, - ), - })) + // An unbounded input may remain pending without ever terminating. Use a + // separate stream that flushes buffered rows at that latency boundary. + if self.input.boundedness().is_unbounded() { + Ok(Box::pin(StreamingFilterStream { + schema: self.schema(), + predicate: Arc::clone(&self.predicate), + input: self.input.execute(partition, context)?, + metrics: FilterExecMetrics::new(&self.metrics, partition), + projection: self.projection.clone(), + batch_coalescer: LimitedBatchCoalescer::new( + self.schema(), + self.batch_size, + self.fetch, + ), + })) + } else { + Ok(Box::pin(FilterExecStream { + schema: self.schema(), + predicate: Arc::clone(&self.predicate), + input: self.input.execute(partition, context)?, + metrics: FilterExecMetrics::new(&self.metrics, partition), + projection: self.projection.clone(), + batch_coalescer: LimitedBatchCoalescer::new( + self.schema(), + self.batch_size, + self.fetch, + ), + })) + } } fn metrics(&self) -> Option { @@ -1332,6 +1348,23 @@ struct FilterExecStream { batch_coalescer: LimitedBatchCoalescer, } +/// A filter stream for unbounded inputs that prioritizes emitting available rows over +/// coalescing them into larger batches. +struct StreamingFilterStream { + /// Output schema after the projection + schema: SchemaRef, + /// The expression to filter on. This expression must evaluate to a boolean value. + predicate: Arc, + /// The input partition to filter. + input: SendableRecordBatchStream, + /// Runtime metrics recording + metrics: FilterExecMetrics, + /// The projection indices of the columns in the input schema + projection: Option, + /// Batch coalescer to combine small batches + batch_coalescer: LimitedBatchCoalescer, +} + /// The metrics for `FilterExec` struct FilterExecMetrics { /// Common metrics for most operators @@ -1394,7 +1427,6 @@ impl Stream for FilterExecStream { ) -> Poll> { let elapsed_compute = self.metrics.baseline_metrics.elapsed_compute().clone(); loop { - // If there is a completed batch ready, return it if let Some(batch) = self.batch_coalescer.next_completed_batch() { self.metrics.selectivity.add_part(batch.num_rows()); let poll = Poll::Ready(Some(Ok(batch))); @@ -1402,83 +1434,151 @@ impl Stream for FilterExecStream { } if self.batch_coalescer.is_finished() { - // If input is done and no batches are ready, return None to signal end of stream. return Poll::Ready(None); } - // Attempt to pull the next batch from the input stream. match ready!(self.input.poll_next_unpin(cx)) { None => { self.batch_coalescer.finish()?; - // Release the input pipeline's resources. let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - // continue draining the coalescer } Some(Ok(batch)) => { let timer = elapsed_compute.timer(); - let status = self.predicate.as_ref() + let status = self + .predicate + .as_ref() .evaluate(&batch) .and_then(|v| v.into_array(batch.num_rows())) .and_then(|array| { - Ok(match self.projection.as_ref() { + Ok(match self.projection.as_ref() { Some(projection) => { let projected_batch = batch.project(projection)?; (array, projected_batch) - }, - None => (array, batch) - }) - }).and_then(|(array, batch)| { - 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)?; - Ok(state) - } - Err(_) => { - internal_err!( - "Cannot create filter_array from non-boolean predicates" - ) } + None => (array, batch), + }) + }) + .and_then(|(array, batch)| 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)?; + self.batch_coalescer.push_batch(batch) } + Err(_) => internal_err!( + "Cannot create filter_array from non-boolean predicates" + ), })?; timer.done(); - match status { - PushBatchStatus::Continue => { - // Keep pushing more batches - } - PushBatchStatus::LimitReached => { - // limit was reached, so stop early - self.batch_coalescer.finish()?; - // Release the input pipeline's resources. - let input_schema = self.input.schema(); - self.input = - Box::pin(EmptyRecordBatchStream::new(input_schema)); - // continue draining the coalescer - } + if status == PushBatchStatus::LimitReached { + self.batch_coalescer.finish()?; + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } } - - // Error case other => return Poll::Ready(other), } } } fn size_hint(&self) -> (usize, Option) { - // Same number of record batches self.input.size_hint() } } + impl RecordBatchStream for FilterExecStream { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) } } +impl Stream for StreamingFilterStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + let elapsed_compute = self.metrics.baseline_metrics.elapsed_compute().clone(); + loop { + if let Some(batch) = self.batch_coalescer.next_completed_batch() { + self.metrics.selectivity.add_part(batch.num_rows()); + let poll = Poll::Ready(Some(Ok(batch))); + return self.metrics.baseline_metrics.record_poll(poll); + } + + if self.batch_coalescer.is_finished() { + return Poll::Ready(None); + } + + match self.input.poll_next_unpin(cx) { + Poll::Pending => { + // An unbounded input may remain pending indefinitely. Flush any rows + // already accepted by the filter so downstream operators can proceed. + if !self.batch_coalescer.is_empty() { + self.batch_coalescer.flush()?; + continue; + } + return Poll::Pending; + } + Poll::Ready(None) => { + self.batch_coalescer.finish()?; + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + Poll::Ready(Some(Ok(batch))) => { + let timer = elapsed_compute.timer(); + let status = self + .predicate + .as_ref() + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows())) + .and_then(|array| { + Ok(match self.projection.as_ref() { + Some(projection) => { + let projected_batch = batch.project(projection)?; + (array, projected_batch) + } + None => (array, batch), + }) + }) + .and_then(|(array, batch)| 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)?; + self.batch_coalescer.push_batch(batch) + } + Err(_) => internal_err!( + "Cannot create filter_array from non-boolean predicates" + ), + })?; + timer.done(); + + if status == PushBatchStatus::LimitReached { + self.batch_coalescer.finish()?; + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + } + Poll::Ready(other) => return Poll::Ready(other), + } + } + } + + fn size_hint(&self) -> (usize, Option) { + self.input.size_hint() + } +} + +impl RecordBatchStream for StreamingFilterStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + /// Return the equals Column-Pairs and Non-equals Column-Pairs #[deprecated( since = "51.0.0", @@ -1540,9 +1640,163 @@ mod tests { use crate::empty::EmptyExec; use crate::expressions::*; use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::stream::RecordBatchStreamAdapter; + use crate::streaming::{PartitionStream, StreamingTableExec}; use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::Int32Array; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + use futures::stream::poll_fn; + use std::fmt::{Debug, Formatter}; + use std::sync::Mutex; + use std::time::Duration; + use tokio::sync::mpsc::{Receiver, Sender, channel}; + + struct ChannelPartition { + schema: SchemaRef, + receiver: Mutex>>>, + } + + impl Debug for ChannelPartition { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChannelPartition").finish_non_exhaustive() + } + } + + impl PartitionStream for ChannelPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let mut receiver = self + .receiver + .lock() + .expect("channel receiver lock poisoned") + .take() + .expect("channel partition can only be executed once"); + let input = poll_fn(move |cx| receiver.poll_recv(cx)); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + input, + )) + } + } + + fn int32_batch(schema: &SchemaRef, values: Vec) -> Result { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(Int32Array::from(values))], + )?) + } + + fn unbounded_int32_filter( + schema: &SchemaRef, + ) -> Result<(Sender>, SendableRecordBatchStream)> { + let (tx, rx) = channel(2); + let source = Arc::new(StreamingTableExec::try_new( + Arc::clone(schema), + vec![Arc::new(ChannelPartition { + schema: Arc::clone(schema), + receiver: Mutex::new(Some(rx)), + })], + None, + vec![], + true, + None, + )?); + let predicate = + binary(col("value", schema)?, Operator::GtEq, lit(10_i32), schema)?; + let filter = Arc::new(FilterExec::try_new(predicate, source)?); + let output = crate::execute_stream(filter, Arc::new(TaskContext::default()))?; + Ok((tx, output)) + } + + #[tokio::test] + async fn unbounded_filter_emits_while_input_is_open() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let (tx, mut output) = unbounded_int32_filter(&schema)?; + + tx.send(Ok(int32_batch(&schema, vec![5, 10, 15])?)) + .await + .expect("query stream should remain open"); + let first = tokio::time::timeout(Duration::from_secs(5), output.next()) + .await + .expect("unbounded filter withheld the first batch") + .expect("unbounded filter ended unexpectedly")?; + assert_eq!(first, int32_batch(&schema, vec![10, 15])?); + + tx.send(Ok(int32_batch(&schema, vec![1, 5])?)) + .await + .expect("query stream should accept a filtered-out batch"); + tx.send(Ok(int32_batch(&schema, vec![20])?)) + .await + .expect("query stream should accept another batch"); + let second = tokio::time::timeout(Duration::from_secs(5), output.next()) + .await + .expect("unbounded filter withheld the second batch") + .expect("unbounded filter ended unexpectedly")?; + assert_eq!(second, int32_batch(&schema, vec![20])?); + + drop(tx); + assert!(output.next().await.is_none()); + Ok(()) + } + + #[tokio::test] + async fn unbounded_filter_coalesces_immediately_ready_batches() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let (tx, mut output) = unbounded_int32_filter(&schema)?; + + tx.send(Ok(int32_batch(&schema, vec![5, 10, 15])?)) + .await + .expect("query stream should remain open"); + tx.send(Ok(int32_batch(&schema, vec![1, 20])?)) + .await + .expect("query stream should remain open"); + + let batch = tokio::time::timeout(Duration::from_secs(5), output.next()) + .await + .expect("unbounded filter withheld ready rows") + .expect("unbounded filter ended unexpectedly")?; + assert_eq!(batch, int32_batch(&schema, vec![10, 15, 20])?); + + drop(tx); + assert!(output.next().await.is_none()); + Ok(()) + } + + #[tokio::test] + async fn bounded_filter_still_coalesces_small_batches() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let input = test::TestMemoryExec::try_new_exec( + &[vec![ + int32_batch(&schema, vec![5, 10, 15])?, + int32_batch(&schema, vec![1, 20])?, + ]], + Arc::clone(&schema), + None, + )?; + let predicate = + binary(col("value", &schema)?, Operator::GtEq, lit(10_i32), &schema)?; + let filter = Arc::new(FilterExec::try_new(predicate, input)?); + + let output = crate::collect(filter, Arc::new(TaskContext::default())).await?; + assert_eq!(output, vec![int32_batch(&schema, vec![10, 15, 20])?]); + Ok(()) + } #[test] fn filter_rejects_zero_batch_size() -> Result<()> {