Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/limit.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! Defines the LIMIT plan
19
20
use std::any::Any;
21
use std::pin::Pin;
22
use std::sync::Arc;
23
use std::task::{Context, Poll};
24
25
use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
26
use super::{
27
    DisplayAs, ExecutionMode, ExecutionPlanProperties, PlanProperties, RecordBatchStream,
28
    SendableRecordBatchStream, Statistics,
29
};
30
use crate::{DisplayFormatType, Distribution, ExecutionPlan, Partitioning};
31
32
use arrow::datatypes::SchemaRef;
33
use arrow::record_batch::RecordBatch;
34
use datafusion_common::{internal_err, Result};
35
use datafusion_execution::TaskContext;
36
37
use futures::stream::{Stream, StreamExt};
38
use log::trace;
39
40
/// Limit execution plan
41
#[derive(Debug)]
42
pub struct GlobalLimitExec {
43
    /// Input execution plan
44
    input: Arc<dyn ExecutionPlan>,
45
    /// Number of rows to skip before fetch
46
    skip: usize,
47
    /// Maximum number of rows to fetch,
48
    /// `None` means fetching all rows
49
    fetch: Option<usize>,
50
    /// Execution metrics
51
    metrics: ExecutionPlanMetricsSet,
52
    cache: PlanProperties,
53
}
54
55
impl GlobalLimitExec {
56
    /// Create a new GlobalLimitExec
57
24
    pub fn new(input: Arc<dyn ExecutionPlan>, skip: usize, fetch: Option<usize>) -> Self {
58
24
        let cache = Self::compute_properties(&input);
59
24
        GlobalLimitExec {
60
24
            input,
61
24
            skip,
62
24
            fetch,
63
24
            metrics: ExecutionPlanMetricsSet::new(),
64
24
            cache,
65
24
        }
66
24
    }
67
68
    /// Input execution plan
69
0
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
70
0
        &self.input
71
0
    }
72
73
    /// Number of rows to skip before fetch
74
0
    pub fn skip(&self) -> usize {
75
0
        self.skip
76
0
    }
77
78
    /// Maximum number of rows to fetch
79
0
    pub fn fetch(&self) -> Option<usize> {
80
0
        self.fetch
81
0
    }
82
83
    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
84
24
    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
85
24
        PlanProperties::new(
86
24
            input.equivalence_properties().clone(), // Equivalence Properties
87
24
            Partitioning::UnknownPartitioning(1),   // Output Partitioning
88
24
            ExecutionMode::Bounded,                 // Execution Mode
89
24
        )
90
24
    }
91
}
92
93
impl DisplayAs for GlobalLimitExec {
94
0
    fn fmt_as(
95
0
        &self,
96
0
        t: DisplayFormatType,
97
0
        f: &mut std::fmt::Formatter,
98
0
    ) -> std::fmt::Result {
99
0
        match t {
100
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
101
0
                write!(
102
0
                    f,
103
0
                    "GlobalLimitExec: skip={}, fetch={}",
104
0
                    self.skip,
105
0
                    self.fetch.map_or("None".to_string(), |x| x.to_string())
106
0
                )
107
0
            }
108
0
        }
109
0
    }
110
}
111
112
impl ExecutionPlan for GlobalLimitExec {
113
0
    fn name(&self) -> &'static str {
114
0
        "GlobalLimitExec"
115
0
    }
116
117
    /// Return a reference to Any that can be used for downcasting
118
0
    fn as_any(&self) -> &dyn Any {
119
0
        self
120
0
    }
121
122
16
    fn properties(&self) -> &PlanProperties {
123
16
        &self.cache
124
16
    }
125
126
0
    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
127
0
        vec![&self.input]
128
0
    }
129
130
0
    fn required_input_distribution(&self) -> Vec<Distribution> {
131
0
        vec![Distribution::SinglePartition]
132
0
    }
133
134
0
    fn maintains_input_order(&self) -> Vec<bool> {
135
0
        vec![true]
136
0
    }
137
138
0
    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
139
0
        vec![false]
140
0
    }
141
142
0
    fn with_new_children(
143
0
        self: Arc<Self>,
144
0
        children: Vec<Arc<dyn ExecutionPlan>>,
145
0
    ) -> Result<Arc<dyn ExecutionPlan>> {
146
0
        Ok(Arc::new(GlobalLimitExec::new(
147
0
            Arc::clone(&children[0]),
148
0
            self.skip,
149
0
            self.fetch,
150
0
        )))
151
0
    }
152
153
8
    fn execute(
154
8
        &self,
155
8
        partition: usize,
156
8
        context: Arc<TaskContext>,
157
8
    ) -> Result<SendableRecordBatchStream> {
158
8
        trace!(
159
0
            "Start GlobalLimitExec::execute for partition: {}",
160
            partition
161
        );
162
        // GlobalLimitExec has a single output partition
163
8
        if 0 != partition {
164
0
            return internal_err!("GlobalLimitExec invalid partition {partition}");
165
8
        }
166
8
167
8
        // GlobalLimitExec requires a single input partition
168
8
        if 1 != self.input.output_partitioning().partition_count() {
169
0
            return internal_err!("GlobalLimitExec requires a single input partition");
170
8
        }
171
8
172
8
        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
173
8
        let stream = self.input.execute(0, context)
?0
;
174
8
        Ok(Box::pin(LimitStream::new(
175
8
            stream,
176
8
            self.skip,
177
8
            self.fetch,
178
8
            baseline_metrics,
179
8
        )))
180
8
    }
181
182
0
    fn metrics(&self) -> Option<MetricsSet> {
183
0
        Some(self.metrics.clone_inner())
184
0
    }
185
186
16
    fn statistics(&self) -> Result<Statistics> {
187
16
        Statistics::with_fetch(
188
16
            self.input.statistics()
?0
,
189
16
            self.schema(),
190
16
            self.fetch,
191
16
            self.skip,
192
            1,
193
        )
194
16
    }
195
196
0
    fn fetch(&self) -> Option<usize> {
197
0
        self.fetch
198
0
    }
199
200
0
    fn supports_limit_pushdown(&self) -> bool {
201
0
        true
202
0
    }
203
}
204
205
/// LocalLimitExec applies a limit to a single partition
206
#[derive(Debug)]
207
pub struct LocalLimitExec {
208
    /// Input execution plan
209
    input: Arc<dyn ExecutionPlan>,
210
    /// Maximum number of rows to return
211
    fetch: usize,
212
    /// Execution metrics
213
    metrics: ExecutionPlanMetricsSet,
214
    cache: PlanProperties,
215
}
216
217
impl LocalLimitExec {
218
    /// Create a new LocalLimitExec partition
219
1
    pub fn new(input: Arc<dyn ExecutionPlan>, fetch: usize) -> Self {
220
1
        let cache = Self::compute_properties(&input);
221
1
        Self {
222
1
            input,
223
1
            fetch,
224
1
            metrics: ExecutionPlanMetricsSet::new(),
225
1
            cache,
226
1
        }
227
1
    }
228
229
    /// Input execution plan
230
0
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
231
0
        &self.input
232
0
    }
233
234
    /// Maximum number of rows to fetch
235
0
    pub fn fetch(&self) -> usize {
236
0
        self.fetch
237
0
    }
238
239
    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
240
1
    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
241
1
        PlanProperties::new(
242
1
            input.equivalence_properties().clone(), // Equivalence Properties
243
1
            input.output_partitioning().clone(),    // Output Partitioning
244
1
            ExecutionMode::Bounded,                 // Execution Mode
245
1
        )
246
1
    }
247
}
248
249
impl DisplayAs for LocalLimitExec {
250
0
    fn fmt_as(
251
0
        &self,
252
0
        t: DisplayFormatType,
253
0
        f: &mut std::fmt::Formatter,
254
0
    ) -> std::fmt::Result {
255
0
        match t {
256
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
257
0
                write!(f, "LocalLimitExec: fetch={}", self.fetch)
258
0
            }
259
0
        }
260
0
    }
261
}
262
263
impl ExecutionPlan for LocalLimitExec {
264
0
    fn name(&self) -> &'static str {
265
0
        "LocalLimitExec"
266
0
    }
267
268
    /// Return a reference to Any that can be used for downcasting
269
0
    fn as_any(&self) -> &dyn Any {
270
0
        self
271
0
    }
272
273
1
    fn properties(&self) -> &PlanProperties {
274
1
        &self.cache
275
1
    }
276
277
0
    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
278
0
        vec![&self.input]
279
0
    }
280
281
0
    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
282
0
        vec![false]
283
0
    }
284
285
0
    fn maintains_input_order(&self) -> Vec<bool> {
286
0
        vec![true]
287
0
    }
288
289
0
    fn with_new_children(
290
0
        self: Arc<Self>,
291
0
        children: Vec<Arc<dyn ExecutionPlan>>,
292
0
    ) -> Result<Arc<dyn ExecutionPlan>> {
293
0
        match children.len() {
294
0
            1 => Ok(Arc::new(LocalLimitExec::new(
295
0
                Arc::clone(&children[0]),
296
0
                self.fetch,
297
0
            ))),
298
0
            _ => internal_err!("LocalLimitExec wrong number of children"),
299
        }
300
0
    }
301
302
0
    fn execute(
303
0
        &self,
304
0
        partition: usize,
305
0
        context: Arc<TaskContext>,
306
0
    ) -> Result<SendableRecordBatchStream> {
307
0
        trace!("Start LocalLimitExec::execute for partition {} of context session_id {} and task_id {:?}", partition, context.session_id(), context.task_id());
308
0
        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
309
0
        let stream = self.input.execute(partition, context)?;
310
0
        Ok(Box::pin(LimitStream::new(
311
0
            stream,
312
0
            0,
313
0
            Some(self.fetch),
314
0
            baseline_metrics,
315
0
        )))
316
0
    }
317
318
0
    fn metrics(&self) -> Option<MetricsSet> {
319
0
        Some(self.metrics.clone_inner())
320
0
    }
321
322
1
    fn statistics(&self) -> Result<Statistics> {
323
1
        Statistics::with_fetch(
324
1
            self.input.statistics()
?0
,
325
1
            self.schema(),
326
1
            Some(self.fetch),
327
            0,
328
            1,
329
        )
330
1
    }
331
332
0
    fn fetch(&self) -> Option<usize> {
333
0
        Some(self.fetch)
334
0
    }
335
336
0
    fn supports_limit_pushdown(&self) -> bool {
337
0
        true
338
0
    }
339
}
340
341
/// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows.
342
pub struct LimitStream {
343
    /// The remaining number of rows to skip
344
    skip: usize,
345
    /// The remaining number of rows to produce
346
    fetch: usize,
347
    /// The input to read from. This is set to None once the limit is
348
    /// reached to enable early termination
349
    input: Option<SendableRecordBatchStream>,
350
    /// Copy of the input schema
351
    schema: SchemaRef,
352
    /// Execution time metrics
353
    baseline_metrics: BaselineMetrics,
354
}
355
356
impl LimitStream {
357
14
    pub fn new(
358
14
        input: SendableRecordBatchStream,
359
14
        skip: usize,
360
14
        fetch: Option<usize>,
361
14
        baseline_metrics: BaselineMetrics,
362
14
    ) -> Self {
363
14
        let schema = input.schema();
364
14
        Self {
365
14
            skip,
366
14
            fetch: fetch.unwrap_or(usize::MAX),
367
14
            input: Some(input),
368
14
            schema,
369
14
            baseline_metrics,
370
14
        }
371
14
    }
372
373
10
    fn poll_and_skip(
374
10
        &mut self,
375
10
        cx: &mut Context<'_>,
376
10
    ) -> Poll<Option<Result<RecordBatch>>> {
377
10
        let input = self.input.as_mut().unwrap();
378
22
        loop {
379
22
            let poll = input.poll_next_unpin(cx);
380
22
            let poll = poll.map_ok(|batch| {
381
14
                if batch.num_rows() <= self.skip {
382
12
                    self.skip -= batch.num_rows();
383
12
                    RecordBatch::new_empty(input.schema())
384
                } else {
385
2
                    let new_batch = batch.slice(self.skip, batch.num_rows() - self.skip);
386
2
                    self.skip = 0;
387
2
                    new_batch
388
                }
389
22
            
}14
);
390
391
14
            match &poll {
392
14
                Poll::Ready(Some(Ok(batch))) => {
393
14
                    if batch.num_rows() > 0 {
394
2
                        break poll;
395
12
                    } else {
396
12
                        // continue to poll input stream
397
12
                    }
398
                }
399
0
                Poll::Ready(Some(Err(_e))) => break poll,
400
3
                Poll::Ready(None) => break poll,
401
5
                Poll::Pending => break poll,
402
            }
403
        }
404
10
    }
405
406
    /// fetches from the batch
407
22
    fn stream_limit(&mut self, batch: RecordBatch) -> Option<RecordBatch> {
408
22
        // records time on drop
409
22
        let _timer = self.baseline_metrics.elapsed_compute().timer();
410
22
        if self.fetch == 0 {
411
0
            self.input = None; // clear input so it can be dropped early
412
0
            None
413
22
        } else if batch.num_rows() < self.fetch {
414
            //
415
13
            self.fetch -= batch.num_rows();
416
13
            Some(batch)
417
9
        } else if batch.num_rows() >= self.fetch {
418
9
            let batch_rows = self.fetch;
419
9
            self.fetch = 0;
420
9
            self.input = None; // clear input so it can be dropped early
421
9
422
9
            // It is guaranteed that batch_rows is <= batch.num_rows
423
9
            Some(batch.slice(0, batch_rows))
424
        } else {
425
0
            unreachable!()
426
        }
427
22
    }
428
}
429
430
impl Stream for LimitStream {
431
    type Item = Result<RecordBatch>;
432
433
44
    fn poll_next(
434
44
        mut self: Pin<&mut Self>,
435
44
        cx: &mut Context<'_>,
436
44
    ) -> Poll<Option<Self::Item>> {
437
44
        let fetch_started = self.skip == 0;
438
44
        let poll = match &mut self.input {
439
35
            Some(input) => {
440
35
                let poll = if fetch_started {
441
25
                    input.poll_next_unpin(cx)
442
                } else {
443
10
                    self.poll_and_skip(cx)
444
                };
445
446
35
                poll.map(|x| 
m27
atch
x22
{
447
22
                    Some(Ok(batch)) => Ok(self.stream_limit(batch)).transpose(),
448
5
                    other => other,
449
35
                
}27
)
450
            }
451
            // input has been cleared
452
9
            None => Poll::Ready(None),
453
        };
454
455
44
        self.baseline_metrics.record_poll(poll)
456
44
    }
457
}
458
459
impl RecordBatchStream for LimitStream {
460
    /// Get the schema
461
0
    fn schema(&self) -> SchemaRef {
462
0
        Arc::clone(&self.schema)
463
0
    }
464
}
465
466
#[cfg(test)]
467
mod tests {
468
    use super::*;
469
    use crate::coalesce_partitions::CoalescePartitionsExec;
470
    use crate::common::collect;
471
    use crate::{common, test};
472
473
    use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
474
    use arrow_array::RecordBatchOptions;
475
    use arrow_schema::Schema;
476
    use datafusion_common::stats::Precision;
477
    use datafusion_physical_expr::expressions::col;
478
    use datafusion_physical_expr::PhysicalExpr;
479
480
    #[tokio::test]
481
1
    async fn limit() -> Result<()> {
482
1
        let task_ctx = Arc::new(TaskContext::default());
483
1
484
1
        let num_partitions = 4;
485
1
        let csv = test::scan_partitioned(num_partitions);
486
1
487
1
        // input should have 4 partitions
488
1
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
489
1
490
1
        let limit =
491
1
            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), 0, Some(7));
492
1
493
1
        // the result should contain 4 batches (one per input partition)
494
1
        let iter = limit.execute(0, task_ctx)
?0
;
495
1
        let batches = common::collect(iter).await
?0
;
496
1
497
1
        // there should be a total of 100 rows
498
1
        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
499
1
        assert_eq!(row_count, 7);
500
1
501
1
        Ok(())
502
1
    }
503
504
    #[tokio::test]
505
1
    async fn limit_early_shutdown() -> Result<()> {
506
1
        let batches = vec![
507
1
            test::make_partition(5),
508
1
            test::make_partition(10),
509
1
            test::make_partition(15),
510
1
            test::make_partition(20),
511
1
            test::make_partition(25),
512
1
        ];
513
1
        let input = test::exec::TestStream::new(batches);
514
1
515
1
        let index = input.index();
516
1
        assert_eq!(index.value(), 0);
517
1
518
1
        // limit of six needs to consume the entire first record batch
519
1
        // (5 rows) and 1 row from the second (1 row)
520
1
        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
521
1
        let limit_stream =
522
1
            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
523
1
        assert_eq!(index.value(), 0);
524
1
525
1
        let results = collect(Box::pin(limit_stream)).
await0
.unwrap();
526
2
        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
527
1
        // Only 6 rows should have been produced
528
1
        assert_eq!(num_rows, 6);
529
1
530
1
        // Only the first two batches should be consumed
531
1
        assert_eq!(index.value(), 2);
532
1
533
1
        Ok(())
534
1
    }
535
536
    #[tokio::test]
537
1
    async fn limit_equals_batch_size() -> Result<()> {
538
1
        let batches = vec![
539
1
            test::make_partition(6),
540
1
            test::make_partition(6),
541
1
            test::make_partition(6),
542
1
        ];
543
1
        let input = test::exec::TestStream::new(batches);
544
1
545
1
        let index = input.index();
546
1
        assert_eq!(index.value(), 0);
547
1
548
1
        // limit of six needs to consume the entire first record batch
549
1
        // (6 rows) and stop immediately
550
1
        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
551
1
        let limit_stream =
552
1
            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
553
1
        assert_eq!(index.value(), 0);
554
1
555
1
        let results = collect(Box::pin(limit_stream)).
await0
.unwrap();
556
1
        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
557
1
        // Only 6 rows should have been produced
558
1
        assert_eq!(num_rows, 6);
559
1
560
1
        // Only the first batch should be consumed
561
1
        assert_eq!(index.value(), 1);
562
1
563
1
        Ok(())
564
1
    }
565
566
    #[tokio::test]
567
1
    async fn limit_no_column() -> Result<()> {
568
1
        let batches = vec![
569
1
            make_batch_no_column(6),
570
1
            make_batch_no_column(6),
571
1
            make_batch_no_column(6),
572
1
        ];
573
1
        let input = test::exec::TestStream::new(batches);
574
1
575
1
        let index = input.index();
576
1
        assert_eq!(index.value(), 0);
577
1
578
1
        // limit of six needs to consume the entire first record batch
579
1
        // (6 rows) and stop immediately
580
1
        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
581
1
        let limit_stream =
582
1
            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
583
1
        assert_eq!(index.value(), 0);
584
1
585
1
        let results = collect(Box::pin(limit_stream)).
await0
.unwrap();
586
1
        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
587
1
        // Only 6 rows should have been produced
588
1
        assert_eq!(num_rows, 6);
589
1
590
1
        // Only the first batch should be consumed
591
1
        assert_eq!(index.value(), 1);
592
1
593
1
        Ok(())
594
1
    }
595
596
    // test cases for "skip"
597
7
    async fn skip_and_fetch(skip: usize, fetch: Option<usize>) -> Result<usize> {
598
7
        let task_ctx = Arc::new(TaskContext::default());
599
7
600
7
        // 4 partitions @ 100 rows apiece
601
7
        let num_partitions = 4;
602
7
        let csv = test::scan_partitioned(num_partitions);
603
7
604
7
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
605
606
7
        let offset =
607
7
            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch);
608
609
        // the result should contain 4 batches (one per input partition)
610
7
        let iter = offset.execute(0, task_ctx)
?0
;
611
7
        let batches = common::collect(iter).await
?0
;
612
10
        
Ok(batches.iter().map(7
|batch| batch.num_rows()
).sum())7
613
7
    }
614
615
    #[tokio::test]
616
1
    async fn skip_none_fetch_none() -> Result<()> {
617
1
        let row_count = skip_and_fetch(0, None).await
?0
;
618
1
        assert_eq!(row_count, 400);
619
1
        Ok(())
620
1
    }
621
622
    #[tokio::test]
623
1
    async fn skip_none_fetch_50() -> Result<()> {
624
1
        let row_count = skip_and_fetch(0, Some(50)).await
?0
;
625
1
        assert_eq!(row_count, 50);
626
1
        Ok(())
627
1
    }
628
629
    #[tokio::test]
630
1
    async fn skip_3_fetch_none() -> Result<()> {
631
1
        // there are total of 400 rows, we skipped 3 rows (offset = 3)
632
1
        let row_count = skip_and_fetch(3, None).await
?0
;
633
1
        assert_eq!(row_count, 397);
634
1
        Ok(())
635
1
    }
636
637
    #[tokio::test]
638
1
    async fn skip_3_fetch_10_stats() -> Result<()> {
639
1
        // there are total of 100 rows, we skipped 3 rows (offset = 3)
640
1
        let row_count = skip_and_fetch(3, Some(10)).await
?0
;
641
1
        assert_eq!(row_count, 10);
642
1
        Ok(())
643
1
    }
644
645
    #[tokio::test]
646
1
    async fn skip_400_fetch_none() -> Result<()> {
647
1
        let row_count = skip_and_fetch(400, None).await
?0
;
648
1
        assert_eq!(row_count, 0);
649
1
        Ok(())
650
1
    }
651
652
    #[tokio::test]
653
1
    async fn skip_400_fetch_1() -> Result<()> {
654
1
        // there are a total of 400 rows
655
1
        let row_count = skip_and_fetch(400, Some(1)).await
?0
;
656
1
        assert_eq!(row_count, 0);
657
1
        Ok(())
658
1
    }
659
660
    #[tokio::test]
661
1
    async fn skip_401_fetch_none() -> Result<()> {
662
1
        // there are total of 400 rows, we skipped 401 rows (offset = 3)
663
1
        let row_count = skip_and_fetch(401, None).await
?0
;
664
1
        assert_eq!(row_count, 0);
665
1
        Ok(())
666
1
    }
667
668
    #[tokio::test]
669
1
    async fn test_row_number_statistics_for_global_limit() -> Result<()> {
670
1
        let row_count = row_number_statistics_for_global_limit(0, Some(10)).
await0
?0
;
671
1
        assert_eq!(row_count, Precision::Exact(10));
672
1
673
1
        let row_count = row_number_statistics_for_global_limit(5, Some(10)).
await0
?0
;
674
1
        assert_eq!(row_count, Precision::Exact(10));
675
1
676
1
        let row_count = row_number_statistics_for_global_limit(400, Some(10)).
await0
?0
;
677
1
        assert_eq!(row_count, Precision::Exact(0));
678
1
679
1
        let row_count = row_number_statistics_for_global_limit(398, Some(10)).
await0
?0
;
680
1
        assert_eq!(row_count, Precision::Exact(2));
681
1
682
1
        let row_count = row_number_statistics_for_global_limit(398, Some(1)).
await0
?0
;
683
1
        assert_eq!(row_count, Precision::Exact(1));
684
1
685
1
        let row_count = row_number_statistics_for_global_limit(398, None).
await0
?0
;
686
1
        assert_eq!(row_count, Precision::Exact(2));
687
1
688
1
        let row_count =
689
1
            row_number_statistics_for_global_limit(0, Some(usize::MAX)).
await0
?0
;
690
1
        assert_eq!(row_count, Precision::Exact(400));
691
1
692
1
        let row_count =
693
1
            row_number_statistics_for_global_limit(398, Some(usize::MAX)).
await0
?0
;
694
1
        assert_eq!(row_count, Precision::Exact(2));
695
1
696
1
        let row_count =
697
1
            row_number_inexact_statistics_for_global_limit(0, Some(10)).
await0
?0
;
698
1
        assert_eq!(row_count, Precision::Inexact(10));
699
1
700
1
        let row_count =
701
1
            row_number_inexact_statistics_for_global_limit(5, Some(10)).
await0
?0
;
702
1
        assert_eq!(row_count, Precision::Inexact(10));
703
1
704
1
        let row_count =
705
1
            row_number_inexact_statistics_for_global_limit(400, Some(10)).
await0
?0
;
706
1
        assert_eq!(row_count, Precision::Exact(0));
707
1
708
1
        let row_count =
709
1
            row_number_inexact_statistics_for_global_limit(398, Some(10)).
await0
?0
;
710
1
        assert_eq!(row_count, Precision::Inexact(2));
711
1
712
1
        let row_count =
713
1
            row_number_inexact_statistics_for_global_limit(398, Some(1)).
await0
?0
;
714
1
        assert_eq!(row_count, Precision::Inexact(1));
715
1
716
1
        let row_count = row_number_inexact_statistics_for_global_limit(398, None).
await0
?0
;
717
1
        assert_eq!(row_count, Precision::Inexact(2));
718
1
719
1
        let row_count =
720
1
            row_number_inexact_statistics_for_global_limit(0, Some(usize::MAX)).
await0
?0
;
721
1
        assert_eq!(row_count, Precision::Inexact(400));
722
1
723
1
        let row_count =
724
1
            row_number_inexact_statistics_for_global_limit(398, Some(usize::MAX)).
await0
?0
;
725
1
        assert_eq!(row_count, Precision::Inexact(2));
726
1
727
1
        Ok(())
728
1
    }
729
730
    #[tokio::test]
731
1
    async fn test_row_number_statistics_for_local_limit() -> Result<()> {
732
1
        let row_count = row_number_statistics_for_local_limit(4, 10).
await0
?0
;
733
1
        assert_eq!(row_count, Precision::Exact(10));
734
1
735
1
        Ok(())
736
1
    }
737
738
8
    async fn row_number_statistics_for_global_limit(
739
8
        skip: usize,
740
8
        fetch: Option<usize>,
741
8
    ) -> Result<Precision<usize>> {
742
8
        let num_partitions = 4;
743
8
        let csv = test::scan_partitioned(num_partitions);
744
8
745
8
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
746
747
8
        let offset =
748
8
            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch);
749
8
750
8
        Ok(offset.statistics()
?0
.num_rows)
751
8
    }
752
753
8
    pub fn build_group_by(
754
8
        input_schema: &SchemaRef,
755
8
        columns: Vec<String>,
756
8
    ) -> PhysicalGroupBy {
757
8
        let mut group_by_expr: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![];
758
8
        for column in columns.iter() {
759
8
            group_by_expr.push((col(column, input_schema).unwrap(), column.to_string()));
760
8
        }
761
8
        PhysicalGroupBy::new_single(group_by_expr.clone())
762
8
    }
763
764
8
    async fn row_number_inexact_statistics_for_global_limit(
765
8
        skip: usize,
766
8
        fetch: Option<usize>,
767
8
    ) -> Result<Precision<usize>> {
768
8
        let num_partitions = 4;
769
8
        let csv = test::scan_partitioned(num_partitions);
770
8
771
8
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
772
773
        // Adding a "GROUP BY i" changes the input stats from Exact to Inexact.
774
8
        let agg = AggregateExec::try_new(
775
8
            AggregateMode::Final,
776
8
            build_group_by(&csv.schema(), vec!["i".to_string()]),
777
8
            vec![],
778
8
            vec![],
779
8
            Arc::clone(&csv),
780
8
            Arc::clone(&csv.schema()),
781
8
        )
?0
;
782
8
        let agg_exec: Arc<dyn ExecutionPlan> = Arc::new(agg);
783
8
784
8
        let offset = GlobalLimitExec::new(
785
8
            Arc::new(CoalescePartitionsExec::new(agg_exec)),
786
8
            skip,
787
8
            fetch,
788
8
        );
789
8
790
8
        Ok(offset.statistics()
?0
.num_rows)
791
8
    }
792
793
1
    async fn row_number_statistics_for_local_limit(
794
1
        num_partitions: usize,
795
1
        fetch: usize,
796
1
    ) -> Result<Precision<usize>> {
797
1
        let csv = test::scan_partitioned(num_partitions);
798
1
799
1
        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
800
801
1
        let offset = LocalLimitExec::new(csv, fetch);
802
1
803
1
        Ok(offset.statistics()
?0
.num_rows)
804
1
    }
805
806
    /// Return a RecordBatch with a single array with row_count sz
807
3
    fn make_batch_no_column(sz: usize) -> RecordBatch {
808
3
        let schema = Arc::new(Schema::empty());
809
3
810
3
        let options = RecordBatchOptions::new().with_row_count(Option::from(sz));
811
3
        RecordBatch::try_new_with_options(schema, vec![], &options).unwrap()
812
3
    }
813
}