| 597 | |
| 598 | #[tokio::test] |
| 599 | async fn limit_equals_batch_size() -> Result<()> { |
| 600 | let batches = vec![ |
| 601 | test::make_partition(6), |
| 602 | test::make_partition(6), |
| 603 | test::make_partition(6), |
| 604 | ]; |
| 605 | let input = test::exec::TestStream::new(batches); |
| 606 | |
| 607 | let index = input.index(); |
| 608 | assert_eq!(index.value(), 0); |
| 609 | |
| 610 | // Limit of six needs to consume the entire first record batch |
| 611 | // (6 rows) and stop immediately |
| 612 | let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); |
| 613 | let limit_stream = |
| 614 | LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics); |
| 615 | assert_eq!(index.value(), 0); |
| 616 | |
| 617 | let results = collect(Box::pin(limit_stream)).await.unwrap(); |
| 618 | let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); |
| 619 | // Only 6 rows should have been produced |
| 620 | assert_eq!(num_rows, 6); |
| 621 | |
| 622 | // Only the first batch should be consumed |
| 623 | assert_eq!(index.value(), 1); |
| 624 | |
| 625 | Ok(()) |
| 626 | } |
| 627 | |
| 628 | #[tokio::test] |
| 629 | async fn limit_no_column() -> Result<()> { |