Fetches from the batch
(&mut self, batch: RecordBatch)
| 467 | |
| 468 | /// Fetches from the batch |
| 469 | fn stream_limit(&mut self, batch: RecordBatch) -> Option<RecordBatch> { |
| 470 | // records time on drop |
| 471 | let _timer = self.baseline_metrics.elapsed_compute().timer(); |
| 472 | if self.fetch == 0 { |
| 473 | self.input = None; // Clear input so it can be dropped early |
| 474 | None |
| 475 | } else if batch.num_rows() < self.fetch { |
| 476 | // |
| 477 | self.fetch -= batch.num_rows(); |
| 478 | Some(batch) |
| 479 | } else if batch.num_rows() >= self.fetch { |
| 480 | let batch_rows = self.fetch; |
| 481 | self.fetch = 0; |
| 482 | self.input = None; // Clear input so it can be dropped early |
| 483 | |
| 484 | // It is guaranteed that batch_rows is <= batch.num_rows |
| 485 | Some(batch.slice(0, batch_rows)) |
| 486 | } else { |
| 487 | unreachable!() |
| 488 | } |
| 489 | } |
| 490 | } |
| 491 | |
| 492 | impl Stream for LimitStream { |
no test coverage detected