datafusion_physical_plan/
limit.rs

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
20use std::pin::Pin;
21use std::sync::Arc;
22use std::task::{Context, Poll};
23
24use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
25use super::{
26    DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream,
27    SendableRecordBatchStream, Statistics,
28};
29use crate::execution_plan::{Boundedness, CardinalityEffect};
30use crate::statistics::{ChildStats, StatisticsArgs};
31use crate::{
32    ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, Partitioning,
33    ReplaceChildrenOptions, validate_child_count,
34};
35
36use arrow::datatypes::SchemaRef;
37use arrow::record_batch::RecordBatch;
38use datafusion_common::tree_node::TreeNodeRecursion;
39use datafusion_common::{Result, assert_eq_or_internal_err};
40use datafusion_execution::TaskContext;
41
42use datafusion_physical_expr::{LexOrdering, PhysicalExpr};
43use futures::stream::{Stream, StreamExt};
44use log::trace;
45
46/// Limit execution plan
47#[derive(Debug, Clone)]
48pub struct GlobalLimitExec {
49    /// Input execution plan
50    input: Arc<dyn ExecutionPlan>,
51    /// Number of rows to skip before fetch
52    skip: usize,
53    /// Maximum number of rows to fetch,
54    /// `None` means fetching all rows
55    fetch: Option<usize>,
56    /// Execution metrics
57    metrics: ExecutionPlanMetricsSet,
58    /// Input ordering that must be preserved so limit pushdown does not change
59    /// which rows are returned.
60    required_ordering: Option<LexOrdering>,
61    cache: Arc<PlanProperties>,
62}
63
64impl GlobalLimitExec {
65    /// Create a new GlobalLimitExec
66    pub fn new(input: Arc<dyn ExecutionPlan>, skip: usize, fetch: Option<usize>) -> Self {
67        let cache = Self::compute_properties(&input);
68        GlobalLimitExec {
69            input,
70            skip,
71            fetch,
72            metrics: ExecutionPlanMetricsSet::new(),
73            required_ordering: None,
74            cache: Arc::new(cache),
75        }
76    }
77
78    /// Input execution plan
79    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
80        &self.input
81    }
82
83    /// Number of rows to skip before fetch
84    pub fn skip(&self) -> usize {
85        self.skip
86    }
87
88    /// Maximum number of rows to fetch
89    pub fn fetch(&self) -> Option<usize> {
90        self.fetch
91    }
92
93    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
94    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
95        PlanProperties::new(
96            input.equivalence_properties().clone(), // Equivalence Properties
97            Partitioning::UnknownPartitioning(1),   // Output Partitioning
98            input.pipeline_behavior(),
99            // Limit operations are always bounded since they output a finite number of rows
100            Boundedness::Bounded,
101        )
102    }
103
104    /// Get the required ordering from limit
105    pub fn required_ordering(&self) -> &Option<LexOrdering> {
106        &self.required_ordering
107    }
108
109    /// Set the required ordering for limit
110    pub fn set_required_ordering(&mut self, required_ordering: Option<LexOrdering>) {
111        self.required_ordering = required_ordering;
112    }
113}
114
115impl DisplayAs for GlobalLimitExec {
116    fn fmt_as(
117        &self,
118        t: DisplayFormatType,
119        f: &mut std::fmt::Formatter,
120    ) -> std::fmt::Result {
121        match t {
122            DisplayFormatType::Default | DisplayFormatType::Verbose => {
123                write!(
124                    f,
125                    "GlobalLimitExec: skip={}, fetch={}",
126                    self.skip,
127                    self.fetch
128                        .map_or_else(|| "None".to_string(), |x| x.to_string())
129                )
130            }
131            DisplayFormatType::TreeRender => {
132                if let Some(fetch) = self.fetch {
133                    writeln!(f, "limit={fetch}")?;
134                }
135                write!(f, "skip={}", self.skip)
136            }
137        }
138    }
139}
140
141impl ExecutionPlan for GlobalLimitExec {
142    fn name(&self) -> &'static str {
143        "GlobalLimitExec"
144    }
145
146    /// Return a reference to Any that can be used for downcasting
147    fn properties(&self) -> &Arc<PlanProperties> {
148        &self.cache
149    }
150
151    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
152        vec![&self.input]
153    }
154
155    fn required_input_distribution(&self) -> Vec<Distribution> {
156        self.input_distribution_requirements().into_per_child()
157    }
158
159    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
160        crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition])
161    }
162
163    fn maintains_input_order(&self) -> Vec<bool> {
164        vec![true]
165    }
166
167    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
168        vec![false]
169    }
170
171    fn apply_expressions(
172        &self,
173        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
174    ) -> Result<TreeNodeRecursion> {
175        Ok(TreeNodeRecursion::Continue)
176    }
177
178    fn replace_children(
179        self: Arc<Self>,
180        mut children: Vec<Arc<dyn ExecutionPlan>>,
181        options: ReplaceChildrenOptions,
182    ) -> Result<Arc<dyn ExecutionPlan>> {
183        validate_child_count!(self, children);
184        match options.children_properties {
185            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
186                input: children.swap_remove(0),
187                metrics: ExecutionPlanMetricsSet::new(),
188                ..Self::clone(&*self)
189            })),
190            ChildrenPropertiesMode::Recompute => {
191                let mut new_limit =
192                    GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch);
193                new_limit.set_required_ordering(self.required_ordering.clone());
194                Ok(Arc::new(new_limit))
195            }
196        }
197    }
198
199    fn with_new_children(
200        self: Arc<Self>,
201        children: Vec<Arc<dyn ExecutionPlan>>,
202    ) -> Result<Arc<dyn ExecutionPlan>> {
203        self.replace_children(
204            children,
205            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
206        )
207    }
208
209    fn with_new_children_and_same_properties(
210        self: Arc<Self>,
211        children: Vec<Arc<dyn ExecutionPlan>>,
212    ) -> Result<Arc<dyn ExecutionPlan>> {
213        self.replace_children(
214            children,
215            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
216        )
217    }
218
219    fn execute(
220        &self,
221        partition: usize,
222        context: Arc<TaskContext>,
223    ) -> Result<SendableRecordBatchStream> {
224        trace!("Start GlobalLimitExec::execute for partition: {partition}");
225        // GlobalLimitExec has a single output partition
226        assert_eq_or_internal_err!(
227            partition,
228            0,
229            "GlobalLimitExec invalid partition {partition}"
230        );
231
232        // GlobalLimitExec requires a single input partition
233        assert_eq_or_internal_err!(
234            self.input.output_partitioning().partition_count(),
235            1,
236            "GlobalLimitExec requires a single input partition"
237        );
238
239        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
240        let stream = self.input.execute(0, context)?;
241        Ok(Box::pin(LimitStream::new(
242            stream,
243            self.skip,
244            self.fetch,
245            baseline_metrics,
246        )))
247    }
248
249    fn metrics(&self) -> Option<MetricsSet> {
250        Some(self.metrics.clone_inner())
251    }
252
253    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
254        vec![ChildStats::At(partition)]
255    }
256
257    fn statistics_from_inputs(
258        &self,
259        input_stats: &[Arc<Statistics>],
260        _args: &StatisticsArgs,
261    ) -> Result<Arc<Statistics>> {
262        let stats = input_stats[0].as_ref().clone();
263        Ok(Arc::new(stats.with_fetch(self.fetch, self.skip, 1)?))
264    }
265
266    fn fetch(&self) -> Option<usize> {
267        self.fetch
268    }
269
270    fn supports_limit_pushdown(&self) -> bool {
271        true
272    }
273
274    #[cfg(feature = "proto")]
275    fn try_to_proto(
276        &self,
277        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
278    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
279        use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto;
280        use datafusion_proto_models::protobuf;
281        let input = ctx.encode_child(self.input())?;
282        let required_ordering = optional_ordering_try_to_proto(
283            self.required_ordering.as_ref(),
284            &ctx.expr_ctx(),
285        )?;
286        Ok(Some(protobuf::PhysicalPlanNode {
287            physical_plan_type: Some(
288                protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new(
289                    protobuf::GlobalLimitExecNode {
290                        input: Some(Box::new(input)),
291                        skip: self.skip() as u32,
292                        fetch: match self.fetch() {
293                            Some(n) => n as i64,
294                            _ => -1, // no limit
295                        },
296                        required_ordering,
297                    },
298                )),
299            ),
300        }))
301    }
302}
303
304#[cfg(feature = "proto")]
305impl GlobalLimitExec {
306    pub fn try_from_proto(
307        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
308        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
309    ) -> Result<Arc<dyn ExecutionPlan>> {
310        use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto;
311        use datafusion_proto_models::protobuf;
312        let limit = crate::expect_plan_variant!(
313            node,
314            protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit,
315            "GlobalLimitExec",
316        );
317        let input = ctx.decode_required_child(
318            limit.input.as_deref(),
319            "GlobalLimitExec",
320            "input",
321        )?;
322        let fetch = if limit.fetch >= 0 {
323            Some(limit.fetch as usize)
324        } else {
325            None
326        };
327        let required_ordering = optional_ordering_try_from_proto(
328            &limit.required_ordering,
329            &ctx.expr_ctx(input.schema().as_ref()),
330        )?;
331        let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch);
332        exec.set_required_ordering(required_ordering);
333        Ok(Arc::new(exec))
334    }
335}
336
337/// LocalLimitExec applies a limit to a single partition
338#[derive(Debug, Clone)]
339pub struct LocalLimitExec {
340    /// Input execution plan
341    input: Arc<dyn ExecutionPlan>,
342    /// Maximum number of rows to return
343    fetch: usize,
344    /// Execution metrics
345    metrics: ExecutionPlanMetricsSet,
346    /// Input ordering that must be preserved so limit pushdown does not change
347    /// which rows are returned.
348    required_ordering: Option<LexOrdering>,
349    cache: Arc<PlanProperties>,
350}
351
352impl LocalLimitExec {
353    /// Create a new LocalLimitExec partition
354    pub fn new(input: Arc<dyn ExecutionPlan>, fetch: usize) -> Self {
355        let cache = Self::compute_properties(&input);
356        Self {
357            input,
358            fetch,
359            metrics: ExecutionPlanMetricsSet::new(),
360            required_ordering: None,
361            cache: Arc::new(cache),
362        }
363    }
364
365    /// Input execution plan
366    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
367        &self.input
368    }
369
370    /// Maximum number of rows to fetch
371    pub fn fetch(&self) -> usize {
372        self.fetch
373    }
374
375    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
376    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
377        PlanProperties::new(
378            input.equivalence_properties().clone(), // Equivalence Properties
379            input.output_partitioning().clone(),    // Output Partitioning
380            input.pipeline_behavior(),
381            // Limit operations are always bounded since they output a finite number of rows
382            Boundedness::Bounded,
383        )
384    }
385
386    /// Get the required ordering from limit
387    pub fn required_ordering(&self) -> &Option<LexOrdering> {
388        &self.required_ordering
389    }
390
391    /// Set the required ordering for limit
392    pub fn set_required_ordering(&mut self, required_ordering: Option<LexOrdering>) {
393        self.required_ordering = required_ordering;
394    }
395}
396
397impl DisplayAs for LocalLimitExec {
398    fn fmt_as(
399        &self,
400        t: DisplayFormatType,
401        f: &mut std::fmt::Formatter,
402    ) -> std::fmt::Result {
403        match t {
404            DisplayFormatType::Default | DisplayFormatType::Verbose => {
405                write!(f, "LocalLimitExec: fetch={}", self.fetch)
406            }
407            DisplayFormatType::TreeRender => {
408                write!(f, "limit={}", self.fetch)
409            }
410        }
411    }
412}
413
414impl ExecutionPlan for LocalLimitExec {
415    fn name(&self) -> &'static str {
416        "LocalLimitExec"
417    }
418
419    /// Return a reference to Any that can be used for downcasting
420    fn properties(&self) -> &Arc<PlanProperties> {
421        &self.cache
422    }
423
424    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
425        vec![&self.input]
426    }
427
428    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
429        vec![false]
430    }
431
432    fn maintains_input_order(&self) -> Vec<bool> {
433        vec![true]
434    }
435
436    fn apply_expressions(
437        &self,
438        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
439    ) -> Result<TreeNodeRecursion> {
440        Ok(TreeNodeRecursion::Continue)
441    }
442
443    fn replace_children(
444        self: Arc<Self>,
445        mut children: Vec<Arc<dyn ExecutionPlan>>,
446        options: ReplaceChildrenOptions,
447    ) -> Result<Arc<dyn ExecutionPlan>> {
448        validate_child_count!(self, children);
449        match options.children_properties {
450            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
451                input: children.swap_remove(0),
452                metrics: ExecutionPlanMetricsSet::new(),
453                ..Self::clone(&*self)
454            })),
455            ChildrenPropertiesMode::Recompute => {
456                let mut new_limit =
457                    LocalLimitExec::new(children.swap_remove(0), self.fetch);
458                new_limit.set_required_ordering(self.required_ordering.clone());
459                Ok(Arc::new(new_limit))
460            }
461        }
462    }
463
464    fn with_new_children(
465        self: Arc<Self>,
466        children: Vec<Arc<dyn ExecutionPlan>>,
467    ) -> Result<Arc<dyn ExecutionPlan>> {
468        self.replace_children(
469            children,
470            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
471        )
472    }
473
474    fn with_new_children_and_same_properties(
475        self: Arc<Self>,
476        children: Vec<Arc<dyn ExecutionPlan>>,
477    ) -> Result<Arc<dyn ExecutionPlan>> {
478        self.replace_children(
479            children,
480            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
481        )
482    }
483
484    fn execute(
485        &self,
486        partition: usize,
487        context: Arc<TaskContext>,
488    ) -> Result<SendableRecordBatchStream> {
489        trace!(
490            "Start LocalLimitExec::execute for partition {} of context session_id {} and task_id {:?}",
491            partition,
492            context.session_id(),
493            context.task_id()
494        );
495        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
496        let stream = self.input.execute(partition, context)?;
497        Ok(Box::pin(LimitStream::new(
498            stream,
499            0,
500            Some(self.fetch),
501            baseline_metrics,
502        )))
503    }
504
505    fn metrics(&self) -> Option<MetricsSet> {
506        Some(self.metrics.clone_inner())
507    }
508
509    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
510        vec![ChildStats::At(partition)]
511    }
512
513    fn statistics_from_inputs(
514        &self,
515        input_stats: &[Arc<Statistics>],
516        _args: &StatisticsArgs,
517    ) -> Result<Arc<Statistics>> {
518        let stats = input_stats[0].as_ref().clone();
519        Ok(Arc::new(stats.with_fetch(Some(self.fetch), 0, 1)?))
520    }
521
522    fn fetch(&self) -> Option<usize> {
523        Some(self.fetch)
524    }
525
526    fn supports_limit_pushdown(&self) -> bool {
527        true
528    }
529
530    fn cardinality_effect(&self) -> CardinalityEffect {
531        CardinalityEffect::LowerEqual
532    }
533
534    #[cfg(feature = "proto")]
535    fn try_to_proto(
536        &self,
537        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
538    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
539        use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto;
540        use datafusion_proto_models::protobuf;
541        let input = ctx.encode_child(self.input())?;
542        let required_ordering = optional_ordering_try_to_proto(
543            self.required_ordering.as_ref(),
544            &ctx.expr_ctx(),
545        )?;
546        Ok(Some(protobuf::PhysicalPlanNode {
547            physical_plan_type: Some(
548                protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new(
549                    protobuf::LocalLimitExecNode {
550                        input: Some(Box::new(input)),
551                        fetch: self.fetch() as u32,
552                        required_ordering,
553                    },
554                )),
555            ),
556        }))
557    }
558}
559
560#[cfg(feature = "proto")]
561impl LocalLimitExec {
562    pub fn try_from_proto(
563        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
564        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
565    ) -> Result<Arc<dyn ExecutionPlan>> {
566        use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto;
567        use datafusion_proto_models::protobuf;
568        let limit = crate::expect_plan_variant!(
569            node,
570            protobuf::physical_plan_node::PhysicalPlanType::LocalLimit,
571            "LocalLimitExec",
572        );
573        let input =
574            ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?;
575        let required_ordering = optional_ordering_try_from_proto(
576            &limit.required_ordering,
577            &ctx.expr_ctx(input.schema().as_ref()),
578        )?;
579        let mut exec = LocalLimitExec::new(input, limit.fetch as usize);
580        exec.set_required_ordering(required_ordering);
581        Ok(Arc::new(exec))
582    }
583}
584
585/// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows.
586pub struct LimitStream {
587    /// The remaining number of rows to skip
588    skip: usize,
589    /// The remaining number of rows to produce
590    fetch: usize,
591    /// The input to read from. This is set to None once the limit is
592    /// reached to enable early termination
593    input: Option<SendableRecordBatchStream>,
594    /// Copy of the input schema
595    schema: SchemaRef,
596    /// Execution time metrics
597    baseline_metrics: BaselineMetrics,
598}
599
600impl LimitStream {
601    pub fn new(
602        input: SendableRecordBatchStream,
603        skip: usize,
604        fetch: Option<usize>,
605        baseline_metrics: BaselineMetrics,
606    ) -> Self {
607        let schema = input.schema();
608        Self {
609            skip,
610            fetch: fetch.unwrap_or(usize::MAX),
611            input: Some(input),
612            schema,
613            baseline_metrics,
614        }
615    }
616
617    fn poll_and_skip(
618        &mut self,
619        cx: &mut Context<'_>,
620    ) -> Poll<Option<Result<RecordBatch>>> {
621        let input = self.input.as_mut().unwrap();
622        loop {
623            let poll = input.poll_next_unpin(cx);
624            let poll = poll.map_ok(|batch| {
625                if batch.num_rows() <= self.skip {
626                    self.skip -= batch.num_rows();
627                    RecordBatch::new_empty(input.schema())
628                } else {
629                    let new_batch = batch.slice(self.skip, batch.num_rows() - self.skip);
630                    self.skip = 0;
631                    new_batch
632                }
633            });
634
635            match &poll {
636                Poll::Ready(Some(Ok(batch))) => {
637                    if batch.num_rows() > 0 {
638                        break poll;
639                    } else {
640                        // Continue to poll input stream
641                    }
642                }
643                Poll::Ready(Some(Err(_e))) => break poll,
644                Poll::Ready(None) => break poll,
645                Poll::Pending => break poll,
646            }
647        }
648    }
649
650    /// Fetches from the batch
651    fn stream_limit(&mut self, batch: RecordBatch) -> Option<RecordBatch> {
652        // records time on drop
653        let _timer = self.baseline_metrics.elapsed_compute().timer();
654        if self.fetch == 0 {
655            self.input = None; // Clear input so it can be dropped early
656            None
657        } else if batch.num_rows() < self.fetch {
658            //
659            self.fetch -= batch.num_rows();
660            Some(batch)
661        } else if batch.num_rows() >= self.fetch {
662            let batch_rows = self.fetch;
663            self.fetch = 0;
664            self.input = None; // Clear input so it can be dropped early
665
666            // It is guaranteed that batch_rows is <= batch.num_rows
667            Some(batch.slice(0, batch_rows))
668        } else {
669            unreachable!()
670        }
671    }
672}
673
674impl Stream for LimitStream {
675    type Item = Result<RecordBatch>;
676
677    fn poll_next(
678        mut self: Pin<&mut Self>,
679        cx: &mut Context<'_>,
680    ) -> Poll<Option<Self::Item>> {
681        let fetch_started = self.skip == 0;
682        let poll = match &mut self.input {
683            Some(input) => {
684                let poll = if fetch_started {
685                    input.poll_next_unpin(cx)
686                } else {
687                    self.poll_and_skip(cx)
688                };
689
690                poll.map(|x| match x {
691                    Some(Ok(batch)) => Ok(self.stream_limit(batch)).transpose(),
692                    other => other,
693                })
694            }
695            // Input has been cleared
696            None => Poll::Ready(None),
697        };
698
699        self.baseline_metrics.record_poll(poll)
700    }
701}
702
703impl RecordBatchStream for LimitStream {
704    /// Get the schema
705    fn schema(&self) -> SchemaRef {
706        Arc::clone(&self.schema)
707    }
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use crate::coalesce_partitions::CoalescePartitionsExec;
714    use crate::common::collect;
715    use crate::statistics::{StatisticsArgs, StatisticsContext};
716    use crate::test;
717
718    use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
719    use arrow::array::RecordBatchOptions;
720    use arrow::compute::SortOptions;
721    use arrow::datatypes::Schema;
722    use datafusion_common::stats::Precision;
723    use datafusion_physical_expr::expressions::col;
724    use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
725
726    #[tokio::test]
727    async fn limit() -> Result<()> {
728        let task_ctx = Arc::new(TaskContext::default());
729
730        let num_partitions = 4;
731        let csv = test::scan_partitioned(num_partitions);
732
733        // Input should have 4 partitions
734        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
735
736        let limit =
737            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), 0, Some(7));
738
739        // The result should contain 4 batches (one per input partition)
740        let iter = limit.execute(0, task_ctx)?;
741        let batches = collect(iter).await?;
742
743        // There should be a total of 100 rows
744        let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
745        assert_eq!(row_count, 7);
746
747        Ok(())
748    }
749
750    #[tokio::test]
751    async fn limit_early_shutdown() -> Result<()> {
752        let batches = vec![
753            test::make_partition(5),
754            test::make_partition(10),
755            test::make_partition(15),
756            test::make_partition(20),
757            test::make_partition(25),
758        ];
759        let input = test::exec::TestStream::new(batches);
760
761        let index = input.index();
762        assert_eq!(index.value(), 0);
763
764        // Limit of six needs to consume the entire first record batch
765        // (5 rows) and 1 row from the second (1 row)
766        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
767        let limit_stream =
768            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
769        assert_eq!(index.value(), 0);
770
771        let results = collect(Box::pin(limit_stream)).await.unwrap();
772        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
773        // Only 6 rows should have been produced
774        assert_eq!(num_rows, 6);
775
776        // Only the first two batches should be consumed
777        assert_eq!(index.value(), 2);
778
779        Ok(())
780    }
781
782    #[tokio::test]
783    async fn limit_equals_batch_size() -> Result<()> {
784        let batches = vec![
785            test::make_partition(6),
786            test::make_partition(6),
787            test::make_partition(6),
788        ];
789        let input = test::exec::TestStream::new(batches);
790
791        let index = input.index();
792        assert_eq!(index.value(), 0);
793
794        // Limit of six needs to consume the entire first record batch
795        // (6 rows) and stop immediately
796        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
797        let limit_stream =
798            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
799        assert_eq!(index.value(), 0);
800
801        let results = collect(Box::pin(limit_stream)).await.unwrap();
802        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
803        // Only 6 rows should have been produced
804        assert_eq!(num_rows, 6);
805
806        // Only the first batch should be consumed
807        assert_eq!(index.value(), 1);
808
809        Ok(())
810    }
811
812    #[tokio::test]
813    async fn limit_no_column() -> Result<()> {
814        let batches = vec![
815            make_batch_no_column(6),
816            make_batch_no_column(6),
817            make_batch_no_column(6),
818        ];
819        let input = test::exec::TestStream::new(batches);
820
821        let index = input.index();
822        assert_eq!(index.value(), 0);
823
824        // Limit of six needs to consume the entire first record batch
825        // (6 rows) and stop immediately
826        let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
827        let limit_stream =
828            LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics);
829        assert_eq!(index.value(), 0);
830
831        let results = collect(Box::pin(limit_stream)).await.unwrap();
832        let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum();
833        // Only 6 rows should have been produced
834        assert_eq!(num_rows, 6);
835
836        // Only the first batch should be consumed
837        assert_eq!(index.value(), 1);
838
839        Ok(())
840    }
841
842    // Test cases for "skip"
843    async fn skip_and_fetch(skip: usize, fetch: Option<usize>) -> Result<usize> {
844        let task_ctx = Arc::new(TaskContext::default());
845
846        // 4 partitions @ 100 rows apiece
847        let num_partitions = 4;
848        let csv = test::scan_partitioned(num_partitions);
849
850        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
851
852        let offset =
853            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch);
854
855        // The result should contain 4 batches (one per input partition)
856        let iter = offset.execute(0, task_ctx)?;
857        let batches = collect(iter).await?;
858        Ok(batches.iter().map(|batch| batch.num_rows()).sum())
859    }
860
861    #[tokio::test]
862    async fn skip_none_fetch_none() -> Result<()> {
863        let row_count = skip_and_fetch(0, None).await?;
864        assert_eq!(row_count, 400);
865        Ok(())
866    }
867
868    #[tokio::test]
869    async fn skip_none_fetch_50() -> Result<()> {
870        let row_count = skip_and_fetch(0, Some(50)).await?;
871        assert_eq!(row_count, 50);
872        Ok(())
873    }
874
875    #[tokio::test]
876    async fn skip_3_fetch_none() -> Result<()> {
877        // There are total of 400 rows, we skipped 3 rows (offset = 3)
878        let row_count = skip_and_fetch(3, None).await?;
879        assert_eq!(row_count, 397);
880        Ok(())
881    }
882
883    #[tokio::test]
884    async fn skip_3_fetch_10_stats() -> Result<()> {
885        // There are total of 100 rows, we skipped 3 rows (offset = 3)
886        let row_count = skip_and_fetch(3, Some(10)).await?;
887        assert_eq!(row_count, 10);
888        Ok(())
889    }
890
891    #[tokio::test]
892    async fn skip_400_fetch_none() -> Result<()> {
893        let row_count = skip_and_fetch(400, None).await?;
894        assert_eq!(row_count, 0);
895        Ok(())
896    }
897
898    #[tokio::test]
899    async fn skip_400_fetch_1() -> Result<()> {
900        // There are a total of 400 rows
901        let row_count = skip_and_fetch(400, Some(1)).await?;
902        assert_eq!(row_count, 0);
903        Ok(())
904    }
905
906    #[tokio::test]
907    async fn skip_401_fetch_none() -> Result<()> {
908        // There are total of 400 rows, we skipped 401 rows (offset = 3)
909        let row_count = skip_and_fetch(401, None).await?;
910        assert_eq!(row_count, 0);
911        Ok(())
912    }
913
914    #[test]
915    fn replace_children_preserves_required_ordering() -> Result<()> {
916        let source = test::scan_partitioned(1);
917        let schema = source.schema();
918        let ordering = LexOrdering::new(vec![PhysicalSortExpr {
919            expr: col("i", &schema)?,
920            options: SortOptions {
921                descending: true,
922                nulls_first: false,
923            },
924        }]);
925
926        let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10));
927        global.set_required_ordering(ordering.clone());
928        let rebuilt = Arc::new(global).replace_children(
929            vec![test::scan_partitioned(1)],
930            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
931        )?;
932        let rebuilt = rebuilt.downcast_ref::<GlobalLimitExec>().unwrap();
933        assert_eq!(rebuilt.required_ordering(), &ordering);
934
935        let mut local = LocalLimitExec::new(source, 10);
936        local.set_required_ordering(ordering.clone());
937        let rebuilt = Arc::new(local).replace_children(
938            vec![test::scan_partitioned(1)],
939            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
940        )?;
941        let rebuilt = rebuilt.downcast_ref::<LocalLimitExec>().unwrap();
942        assert_eq!(rebuilt.required_ordering(), &ordering);
943
944        Ok(())
945    }
946
947    #[test]
948    fn test_row_number_statistics_for_global_limit() -> Result<()> {
949        let row_count = row_number_statistics_for_global_limit(0, Some(10))?;
950        assert_eq!(row_count, Precision::Exact(10));
951
952        let row_count = row_number_statistics_for_global_limit(5, Some(10))?;
953        assert_eq!(row_count, Precision::Exact(10));
954
955        let row_count = row_number_statistics_for_global_limit(400, Some(10))?;
956        assert_eq!(row_count, Precision::Exact(0));
957
958        let row_count = row_number_statistics_for_global_limit(398, Some(10))?;
959        assert_eq!(row_count, Precision::Exact(2));
960
961        let row_count = row_number_statistics_for_global_limit(398, Some(1))?;
962        assert_eq!(row_count, Precision::Exact(1));
963
964        let row_count = row_number_statistics_for_global_limit(398, None)?;
965        assert_eq!(row_count, Precision::Exact(2));
966
967        let row_count = row_number_statistics_for_global_limit(0, Some(usize::MAX))?;
968        assert_eq!(row_count, Precision::Exact(400));
969
970        let row_count = row_number_statistics_for_global_limit(398, Some(usize::MAX))?;
971        assert_eq!(row_count, Precision::Exact(2));
972
973        let row_count = row_number_inexact_statistics_for_global_limit(0, Some(10))?;
974        assert_eq!(row_count, Precision::Inexact(10));
975
976        let row_count = row_number_inexact_statistics_for_global_limit(5, Some(10))?;
977        assert_eq!(row_count, Precision::Inexact(10));
978
979        // Input was Inexact, so an `nr <= skip` outcome must remain Inexact:
980        // the inexact estimate could be wrong, so we cannot promote 0 to
981        // Exact.
982        let row_count = row_number_inexact_statistics_for_global_limit(400, Some(10))?;
983        assert_eq!(row_count, Precision::Inexact(0));
984
985        let row_count = row_number_inexact_statistics_for_global_limit(398, Some(10))?;
986        assert_eq!(row_count, Precision::Inexact(2));
987
988        let row_count = row_number_inexact_statistics_for_global_limit(398, Some(1))?;
989        assert_eq!(row_count, Precision::Inexact(1));
990
991        let row_count = row_number_inexact_statistics_for_global_limit(398, None)?;
992        assert_eq!(row_count, Precision::Inexact(2));
993
994        let row_count =
995            row_number_inexact_statistics_for_global_limit(0, Some(usize::MAX))?;
996        assert_eq!(row_count, Precision::Inexact(400));
997
998        let row_count =
999            row_number_inexact_statistics_for_global_limit(398, Some(usize::MAX))?;
1000        assert_eq!(row_count, Precision::Inexact(2));
1001
1002        Ok(())
1003    }
1004
1005    #[test]
1006    fn test_row_number_statistics_for_local_limit() -> Result<()> {
1007        let row_count = row_number_statistics_for_local_limit(4, 10)?;
1008        assert_eq!(row_count, Precision::Exact(10));
1009
1010        Ok(())
1011    }
1012
1013    fn row_number_statistics_for_global_limit(
1014        skip: usize,
1015        fetch: Option<usize>,
1016    ) -> Result<Precision<usize>> {
1017        let num_partitions = 4;
1018        let csv = test::scan_partitioned(num_partitions);
1019
1020        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
1021
1022        let offset =
1023            GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch);
1024
1025        Ok(StatisticsContext::new()
1026            .compute(&offset, &StatisticsArgs::new())?
1027            .num_rows)
1028    }
1029
1030    pub fn build_group_by(
1031        input_schema: &SchemaRef,
1032        columns: Vec<String>,
1033    ) -> PhysicalGroupBy {
1034        let mut group_by_expr: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![];
1035        for column in columns.iter() {
1036            group_by_expr.push((col(column, input_schema).unwrap(), column.to_string()));
1037        }
1038        PhysicalGroupBy::new_single(group_by_expr.clone())
1039    }
1040
1041    fn row_number_inexact_statistics_for_global_limit(
1042        skip: usize,
1043        fetch: Option<usize>,
1044    ) -> Result<Precision<usize>> {
1045        let num_partitions = 4;
1046        let csv = test::scan_partitioned(num_partitions);
1047
1048        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
1049
1050        // Adding a "GROUP BY i" changes the input stats from Exact to Inexact.
1051        let agg = AggregateExec::try_new(
1052            AggregateMode::Final,
1053            build_group_by(&csv.schema(), vec!["i".to_string()]),
1054            vec![],
1055            vec![],
1056            Arc::clone(&csv),
1057            Arc::clone(&csv.schema()),
1058        )?;
1059        let agg_exec: Arc<dyn ExecutionPlan> = Arc::new(agg);
1060
1061        let offset = GlobalLimitExec::new(
1062            Arc::new(CoalescePartitionsExec::new(agg_exec)),
1063            skip,
1064            fetch,
1065        );
1066
1067        Ok(StatisticsContext::new()
1068            .compute(&offset, &StatisticsArgs::new())?
1069            .num_rows)
1070    }
1071
1072    fn row_number_statistics_for_local_limit(
1073        num_partitions: usize,
1074        fetch: usize,
1075    ) -> Result<Precision<usize>> {
1076        let csv = test::scan_partitioned(num_partitions);
1077
1078        assert_eq!(csv.output_partitioning().partition_count(), num_partitions);
1079
1080        let offset = LocalLimitExec::new(csv, fetch);
1081
1082        Ok(StatisticsContext::new()
1083            .compute(&offset, &StatisticsArgs::new())?
1084            .num_rows)
1085    }
1086
1087    /// Return a RecordBatch with a single array with row_count sz
1088    fn make_batch_no_column(sz: usize) -> RecordBatch {
1089        let schema = Arc::new(Schema::empty());
1090
1091        let options = RecordBatchOptions::new().with_row_count(Option::from(sz));
1092        RecordBatch::try_new_with_options(schema, vec![], &options).unwrap()
1093    }
1094}
datafusion_physical_plan/limit.rs

datafusion_physical_plan/
limit.rs
