alamb commented on a change in pull request #9926:
URL: https://github.com/apache/arrow/pull/9926#discussion_r608654403
##########
File path: rust/datafusion/src/physical_plan/limit.rs
##########
@@ -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,
- input: SendableRecordBatchStream,
- // the current count
+ /// The input to read from. This is set to None once the limit is
+ /// reached to enable early termination
+ 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();
Review comment:
The actual code change for this PR is very small -- the rest of the
changes are related to writing a proper test for it
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
