ARROW-12254: [Rust][DataFusion] Stop polling limit input once limit is reached #9926
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -200,30 +200,39 @@ pub fn truncate_batch(batch: &RecordBatch, n: usize) -> RecordBatch { | |
|
|
||
| /// A Limit stream limits the stream to up to `limit` rows. | ||
| struct LimitStream { | ||
| /// The maximum number of rows to produce | ||
| limit: usize, | ||
| /// The input to read from. This is set to None once the limit is | ||
| /// reached to enable early termination | ||
alamb marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| input: Option<SendableRecordBatchStream>, | ||
| /// Copy of the input schema | ||
| schema: SchemaRef, | ||
| // the current number of rows which have been produced | ||
| current_len: usize, | ||
| } | ||
|
|
||
| impl LimitStream { | ||
| fn new(input: SendableRecordBatchStream, limit: usize) -> Self { | ||
| let schema = input.schema(); | ||
|
||
| Self { | ||
| limit, | ||
| input: Some(input), | ||
| schema, | ||
| current_len: 0, | ||
| } | ||
| } | ||
|
|
||
| fn stream_limit(&mut self, batch: RecordBatch) -> Option<RecordBatch> { | ||
| if self.current_len == self.limit { | ||
| self.input = None; // clear input so it can be dropped early | ||
| None | ||
| } else if self.current_len + batch.num_rows() <= self.limit { | ||
| self.current_len += batch.num_rows(); | ||
| Some(batch) | ||
| } else { | ||
| let batch_rows = self.limit - self.current_len; | ||
| self.current_len = self.limit; | ||
| self.input = None; // clear input so it can be dropped early | ||
| Some(truncate_batch(&batch, batch_rows)) | ||
| } | ||
| } | ||
|
|
@@ -236,23 +245,29 @@ impl Stream for LimitStream { | |
| mut self: Pin<&mut Self>, | ||
| cx: &mut Context<'_>, | ||
| ) -> Poll<Option<Self::Item>> { | ||
| match &mut self.input { | ||
|
||
| Some(input) => input.poll_next_unpin(cx).map(|x| match x { | ||
| Some(Ok(batch)) => Ok(self.stream_limit(batch)).transpose(), | ||
| other => other, | ||
| }), | ||
| // input has been cleared | ||
| None => Poll::Ready(None), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl RecordBatchStream for LimitStream { | ||
| /// Get the schema | ||
| fn schema(&self) -> SchemaRef { | ||
| self.schema.clone() | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
|
|
||
| use common::collect; | ||
|
|
||
| use super::*; | ||
| use crate::physical_plan::common; | ||
| use crate::physical_plan::csv::{CsvExec, CsvReadOptions}; | ||
|
|
@@ -290,4 +305,34 @@ mod tests { | |
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn limit_early_shutdown() -> Result<()> { | ||
| let batches = vec![ | ||
| test::make_partition(5), | ||
| test::make_partition(10), | ||
| test::make_partition(15), | ||
| test::make_partition(20), | ||
| test::make_partition(25), | ||
| ]; | ||
| let input = test::exec::TestStream::new(batches); | ||
|
|
||
| let index = input.index(); | ||
| assert_eq!(index.value(), 0); | ||
|
|
||
| // limit of six needs to consume the entire first record batch | ||
| // (5 rows) and 1 row from the second (1 row) | ||
| let limit_stream = LimitStream::new(Box::pin(input), 6); | ||
| assert_eq!(index.value(), 0); | ||
|
|
||
| let results = collect(Box::pin(limit_stream)).await.unwrap(); | ||
| let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); | ||
| // Only 6 rows should have been produced | ||
| assert_eq!(num_rows, 6); | ||
|
|
||
| // Only the first two batches should be consumed | ||
| assert_eq!(index.value(), 2); | ||
|
|
||
| Ok(()) | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this necessary? It can be based on
self.current_len == self.limitor otherwise a boolean likelimit_exhausted?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Snap - #9926 (comment)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
haha 👍 😆