query/optimizer/
global_limit.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use datafusion::config::ConfigOptions;
18use datafusion::physical_optimizer::PhysicalOptimizerRule;
19use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
20use datafusion::physical_plan::filter::FilterExec;
21use datafusion::physical_plan::limit::GlobalLimitExec;
22use datafusion::physical_plan::repartition::RepartitionExec;
23use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
24use datafusion::physical_plan::{
25    ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, ReplaceChildrenOptions,
26};
27use datafusion_common::Result as DfResult;
28use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning};
29
30use crate::dist_plan::MergeSortExec;
31
32#[derive(Debug)]
33pub struct EnsureGlobalLimitForFetch;
34
35impl PhysicalOptimizerRule for EnsureGlobalLimitForFetch {
36    fn optimize(
37        &self,
38        plan: Arc<dyn ExecutionPlan>,
39        _config: &ConfigOptions,
40    ) -> DfResult<Arc<dyn ExecutionPlan>> {
41        Self::optimize_plan(plan, ParentContext::default())
42    }
43
44    fn name(&self) -> &str {
45        "EnsureGlobalLimitForFetch"
46    }
47
48    fn schema_check(&self) -> bool {
49        true
50    }
51}
52
53impl EnsureGlobalLimitForFetch {
54    fn optimize_plan(
55        plan: Arc<dyn ExecutionPlan>,
56        parent: ParentContext,
57    ) -> DfResult<Arc<dyn ExecutionPlan>> {
58        let children = plan.children();
59        let plan = if children.is_empty() {
60            plan
61        } else {
62            let required_input_distribution = plan.input_distribution_requirements();
63            let required_input_ordering = plan.required_input_ordering();
64            let maintains_input_order = plan.maintains_input_order();
65            let child_parent = ParentContext {
66                global_fetch: provided_global_fetch(&plan),
67                required_ordering: None,
68                required_distribution: Distribution::UnspecifiedDistribution,
69                partitioning_to_restore: None,
70                preserve_hash_partitioning: false,
71            };
72            let children = children
73                .into_iter()
74                .enumerate()
75                .map(|(idx, child)| {
76                    let required_distribution = required_input_distribution
77                        .child_distribution(idx)
78                        .cloned()
79                        .unwrap_or(Distribution::UnspecifiedDistribution);
80                    let partitioning_to_restore =
81                        partitioning_to_restore_for(child, &required_distribution)
82                            .or_else(|| inherited_partitioning_to_restore(&plan, child, &parent));
83                    let preserve_hash_partitioning = partitioning_to_restore.is_some();
84                    let required_ordering = required_input_ordering
85                        .get(idx)
86                        .cloned()
87                        .unwrap_or(None)
88                        .or_else(|| {
89                            maintains_input_order
90                                .get(idx)
91                                .copied()
92                                .unwrap_or(false)
93                                .then(|| parent.required_ordering.clone())
94                                .flatten()
95                        });
96                    let parent = ParentContext {
97                        required_ordering,
98                        required_distribution,
99                        partitioning_to_restore,
100                        preserve_hash_partitioning,
101                        ..child_parent.clone()
102                    };
103                    Self::optimize_plan(Arc::clone(child), parent)
104                })
105                .collect::<DfResult<Vec<_>>>()?;
106            plan.replace_children(
107                children,
108                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
109            )?
110        };
111
112        let Some(fetch) = plan.fetch() else {
113            return Ok(plan);
114        };
115
116        if parent
117            .global_fetch
118            .is_some_and(|parent_fetch| parent_fetch <= fetch)
119            || !plan.is::<FilterExec>()
120            || plan.output_partitioning().partition_count() <= 1
121        {
122            return Ok(plan);
123        }
124
125        add_global_fetch(
126            plan,
127            fetch,
128            parent.required_ordering,
129            parent.partitioning_to_restore,
130        )
131    }
132}
133
134#[derive(Clone)]
135struct ParentContext {
136    global_fetch: Option<usize>,
137    required_ordering: Option<OrderingRequirements>,
138    required_distribution: Distribution,
139    partitioning_to_restore: Option<Partitioning>,
140    preserve_hash_partitioning: bool,
141}
142
143impl Default for ParentContext {
144    fn default() -> Self {
145        Self {
146            global_fetch: None,
147            required_ordering: None,
148            required_distribution: Distribution::UnspecifiedDistribution,
149            partitioning_to_restore: None,
150            preserve_hash_partitioning: false,
151        }
152    }
153}
154
155fn provided_global_fetch(plan: &Arc<dyn ExecutionPlan>) -> Option<usize> {
156    let fetch = plan.fetch()?;
157    (plan.is::<GlobalLimitExec>()
158        || plan.is::<CoalescePartitionsExec>()
159        || plan.is::<SortPreservingMergeExec>()
160        || plan.is::<MergeSortExec>())
161    .then_some(fetch)
162}
163
164fn add_global_fetch(
165    plan: Arc<dyn ExecutionPlan>,
166    fetch: usize,
167    required_ordering: Option<OrderingRequirements>,
168    partitioning_to_restore: Option<Partitioning>,
169) -> DfResult<Arc<dyn ExecutionPlan>> {
170    let plan = if required_ordering.is_some()
171        && let Some(ordering) = plan.output_ordering().cloned()
172    {
173        Arc::new(SortPreservingMergeExec::new(ordering, plan).with_fetch(Some(fetch)))
174            as Arc<dyn ExecutionPlan>
175    } else {
176        Arc::new(CoalescePartitionsExec::new(plan).with_fetch(Some(fetch)))
177            as Arc<dyn ExecutionPlan>
178    };
179
180    restore_required_partitioning(plan, partitioning_to_restore)
181}
182
183fn restore_required_partitioning(
184    plan: Arc<dyn ExecutionPlan>,
185    partitioning_to_restore: Option<Partitioning>,
186) -> DfResult<Arc<dyn ExecutionPlan>> {
187    let Some(partitioning) = partitioning_to_restore else {
188        return Ok(plan);
189    };
190
191    if partitioning.partition_count() <= 1 || !matches!(&partitioning, Partitioning::Hash(_, _)) {
192        return Ok(plan);
193    }
194
195    Ok(Arc::new(
196        RepartitionExec::try_new(plan, partitioning)?.with_preserve_order(),
197    ))
198}
199
200fn partitioning_to_restore_for(
201    child: &Arc<dyn ExecutionPlan>,
202    required_distribution: &Distribution,
203) -> Option<Partitioning> {
204    if !matches!(required_distribution, Distribution::KeyPartitioned(_))
205        || child.output_partitioning().partition_count() <= 1
206    {
207        return None;
208    }
209
210    if child
211        .output_partitioning()
212        .satisfaction(required_distribution, child.equivalence_properties(), false)
213        .is_satisfied()
214    {
215        Some(child.output_partitioning().clone())
216    } else {
217        Some(
218            required_distribution
219                .clone()
220                .create_partitioning(child.output_partitioning().partition_count()),
221        )
222    }
223}
224
225fn inherited_partitioning_to_restore(
226    plan: &Arc<dyn ExecutionPlan>,
227    child: &Arc<dyn ExecutionPlan>,
228    parent: &ParentContext,
229) -> Option<Partitioning> {
230    if child.output_partitioning().partition_count() <= 1
231        || !matches!(child.output_partitioning(), Partitioning::Hash(_, _))
232        || !matches!(plan.output_partitioning(), Partitioning::Hash(_, _))
233        || plan.output_partitioning().partition_count()
234            != child.output_partitioning().partition_count()
235    {
236        return None;
237    }
238
239    let satisfies_parent_distribution = matches!(
240        parent.required_distribution,
241        Distribution::KeyPartitioned(_)
242    ) && plan
243        .output_partitioning()
244        .satisfaction(
245            &parent.required_distribution,
246            plan.equivalence_properties(),
247            false,
248        )
249        .is_satisfied();
250
251    (satisfies_parent_distribution || parent.preserve_hash_partitioning)
252        .then(|| child.output_partitioning().clone())
253}
254
255#[cfg(test)]
256mod tests {
257    use datafusion::arrow::array::Int32Array;
258    use datafusion::arrow::compute::SortOptions;
259    use datafusion::arrow::datatypes::{DataType, Field, Schema};
260    use datafusion::arrow::record_batch::RecordBatch;
261    use datafusion::physical_expr::expressions::{col, lit};
262    use datafusion::physical_plan::filter::FilterExecBuilder;
263    use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
264    use datafusion::physical_plan::limit::GlobalLimitExec;
265    use datafusion::physical_plan::projection::ProjectionExec;
266    use datafusion::physical_plan::repartition::RepartitionExec;
267    use datafusion::physical_plan::test::TestMemoryExec;
268    use datafusion_common::{JoinType, NullEquality};
269    use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr};
270
271    use super::*;
272
273    #[test]
274    fn adds_global_limit_for_multi_partition_filter_fetch() {
275        let filter = filter_fetch(unordered_input(), 1);
276
277        let optimized =
278            EnsureGlobalLimitForFetch::optimize_plan(filter, ParentContext::default()).unwrap();
279
280        assert!(optimized.is::<CoalescePartitionsExec>());
281        assert_eq!(optimized.fetch(), Some(1));
282        assert_eq!(optimized.output_partitioning().partition_count(), 1);
283    }
284
285    #[test]
286    fn still_visits_subtree_under_global_limit() {
287        let filter = filter_fetch(unordered_input(), 5);
288        let projection = Arc::new(
289            ProjectionExec::try_new(
290                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
291                filter,
292            )
293            .unwrap(),
294        );
295        let limit =
296            Arc::new(GlobalLimitExec::new(projection, 0, Some(10))) as Arc<dyn ExecutionPlan>;
297
298        let optimized =
299            EnsureGlobalLimitForFetch::optimize_plan(limit, ParentContext::default()).unwrap();
300        let projection = optimized.children()[0];
301        let coalesce = projection.children()[0];
302
303        assert!(coalesce.is::<CoalescePartitionsExec>());
304        assert_eq!(coalesce.fetch(), Some(5));
305    }
306
307    #[test]
308    fn keeps_filter_under_parent_global_fetch() {
309        let (input, ordering) = ordered_input();
310        let filter = filter_fetch(input, 1);
311        let merge = Arc::new(SortPreservingMergeExec::new(ordering, filter).with_fetch(Some(1)))
312            as Arc<dyn ExecutionPlan>;
313
314        let optimized =
315            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
316        let child = optimized.children()[0];
317
318        assert!(optimized.is::<SortPreservingMergeExec>());
319        assert!(child.is::<FilterExec>());
320    }
321
322    #[test]
323    fn adds_tighter_global_fetch_under_looser_parent_fetch() {
324        let (input, ordering) = ordered_input();
325        let filter = filter_fetch(input, 5);
326        let merge = Arc::new(SortPreservingMergeExec::new(ordering, filter).with_fetch(Some(10)))
327            as Arc<dyn ExecutionPlan>;
328
329        let optimized =
330            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
331        let child = optimized.children()[0];
332
333        assert!(optimized.is::<SortPreservingMergeExec>());
334        assert!(child.is::<SortPreservingMergeExec>());
335        assert_eq!(child.fetch(), Some(5));
336    }
337
338    #[test]
339    fn keeps_filter_under_parent_merge_sort_fetch() {
340        let (input, ordering) = ordered_input();
341        let filter = filter_fetch(input, 1);
342        let merge = merge_sort_fetch(ordering, filter, 1);
343
344        let optimized =
345            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
346        let child = optimized.children()[0];
347
348        assert!(optimized.is::<MergeSortExec>());
349        assert!(child.is::<FilterExec>());
350    }
351
352    #[test]
353    fn adds_tighter_global_fetch_under_looser_merge_sort_fetch() {
354        let (input, ordering) = ordered_input();
355        let filter = filter_fetch(input, 5);
356        let merge = merge_sort_fetch(ordering, filter, 10);
357
358        let optimized =
359            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
360        let child = optimized.children()[0];
361
362        assert!(optimized.is::<MergeSortExec>());
363        assert!(child.is::<SortPreservingMergeExec>());
364        assert_eq!(child.fetch(), Some(5));
365        assert!(child.children()[0].is::<FilterExec>());
366    }
367
368    #[test]
369    fn preserves_parent_ordering_requirement() {
370        let (input, ordering) = ordered_input();
371        let filter = filter_fetch(input, 1);
372        let merge =
373            Arc::new(SortPreservingMergeExec::new(ordering, filter)) as Arc<dyn ExecutionPlan>;
374
375        let optimized =
376            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
377        let child = optimized.children()[0];
378
379        assert!(optimized.is::<SortPreservingMergeExec>());
380        assert!(child.is::<SortPreservingMergeExec>());
381        assert_eq!(child.fetch(), Some(1));
382    }
383
384    #[test]
385    fn uses_child_output_ordering_for_merge() {
386        let schema = schema();
387        let required_ordering = ordering(schema.as_ref(), false);
388        let actual_ordering = ordering(schema.as_ref(), true);
389        let batch = batch(schema.clone());
390        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
391        let input = TestMemoryExec::try_new(&partitions, schema, None)
392            .unwrap()
393            .try_with_sort_information(vec![actual_ordering.clone()])
394            .unwrap();
395        let filter = filter_fetch(Arc::new(input), 1);
396
397        let optimized = add_global_fetch(
398            filter,
399            1,
400            Some(OrderingRequirements::from(required_ordering)),
401            None,
402        )
403        .unwrap();
404        let merge = optimized.downcast_ref::<SortPreservingMergeExec>().unwrap();
405
406        assert_eq!(merge.expr(), &actual_ordering);
407    }
408
409    #[test]
410    fn preserves_inherited_ordering_requirement_through_projection() {
411        let (input, ordering) = ordered_input();
412        let filter = filter_fetch(input, 1);
413        let projection = Arc::new(
414            ProjectionExec::try_new(
415                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
416                filter,
417            )
418            .unwrap(),
419        );
420        let merge =
421            Arc::new(SortPreservingMergeExec::new(ordering, projection)) as Arc<dyn ExecutionPlan>;
422
423        let optimized =
424            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
425        let projection = optimized.children()[0];
426        let child = projection.children()[0];
427
428        assert!(optimized.is::<SortPreservingMergeExec>());
429        assert!(projection.is::<ProjectionExec>());
430        assert!(child.is::<SortPreservingMergeExec>());
431        assert_eq!(child.fetch(), Some(1));
432    }
433
434    #[test]
435    fn restores_parent_hash_distribution_after_global_fetch() {
436        let left = filter_fetch(hash_repartition(unordered_input()), 1);
437        let right = hash_repartition(unordered_input());
438        let on = vec![(
439            col("a", left.schema().as_ref()).unwrap(),
440            col("a", right.schema().as_ref()).unwrap(),
441        )];
442        let join = Arc::new(
443            HashJoinExec::try_new(
444                left,
445                right,
446                on,
447                None,
448                &JoinType::Inner,
449                None,
450                PartitionMode::Partitioned,
451                NullEquality::NullEqualsNothing,
452                false,
453            )
454            .unwrap(),
455        ) as Arc<dyn ExecutionPlan>;
456
457        let optimized =
458            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
459        let left = optimized.children()[0];
460        let repartition = left.downcast_ref::<RepartitionExec>().unwrap();
461
462        assert!(matches!(
463            repartition.partitioning(),
464            Partitioning::Hash(_, 3)
465        ));
466        assert!(repartition.input().is::<CoalescePartitionsExec>());
467        assert_eq!(repartition.input().fetch(), Some(1));
468    }
469
470    #[test]
471    fn restores_inherited_hash_distribution_through_projection() {
472        let filter = filter_fetch(hash_repartition(unordered_input()), 1);
473        let projection = Arc::new(
474            ProjectionExec::try_new(
475                vec![(col("a", filter.schema().as_ref()).unwrap(), "a".to_string())],
476                filter,
477            )
478            .unwrap(),
479        ) as Arc<dyn ExecutionPlan>;
480        let right = hash_repartition(unordered_input());
481        let on = vec![(
482            col("a", projection.schema().as_ref()).unwrap(),
483            col("a", right.schema().as_ref()).unwrap(),
484        )];
485        let join = Arc::new(
486            HashJoinExec::try_new(
487                projection,
488                right,
489                on,
490                None,
491                &JoinType::Inner,
492                None,
493                PartitionMode::Partitioned,
494                NullEquality::NullEqualsNothing,
495                false,
496            )
497            .unwrap(),
498        ) as Arc<dyn ExecutionPlan>;
499
500        let optimized =
501            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
502        let projection = optimized.children()[0];
503        let repartition = projection.children()[0]
504            .downcast_ref::<RepartitionExec>()
505            .unwrap();
506
507        assert!(projection.is::<ProjectionExec>());
508        assert!(matches!(
509            repartition.partitioning(),
510            Partitioning::Hash(_, 3)
511        ));
512        assert!(repartition.input().is::<CoalescePartitionsExec>());
513        assert_eq!(repartition.input().fetch(), Some(1));
514    }
515
516    #[test]
517    fn restores_inherited_hash_distribution_through_multiple_projections() {
518        let filter = filter_fetch(hash_repartition(unordered_input()), 1);
519        let projection = project_a(filter);
520        let projection = project_a(projection);
521        let right = hash_repartition(unordered_input());
522        let on = vec![(
523            col("a", projection.schema().as_ref()).unwrap(),
524            col("a", right.schema().as_ref()).unwrap(),
525        )];
526        let join = Arc::new(
527            HashJoinExec::try_new(
528                projection,
529                right,
530                on,
531                None,
532                &JoinType::Inner,
533                None,
534                PartitionMode::Partitioned,
535                NullEquality::NullEqualsNothing,
536                false,
537            )
538            .unwrap(),
539        ) as Arc<dyn ExecutionPlan>;
540
541        let optimized =
542            EnsureGlobalLimitForFetch::optimize_plan(join, ParentContext::default()).unwrap();
543        let outer_projection = optimized.children()[0];
544        let inner_projection = outer_projection.children()[0];
545        let repartition = inner_projection.children()[0]
546            .downcast_ref::<RepartitionExec>()
547            .unwrap();
548
549        assert!(outer_projection.is::<ProjectionExec>());
550        assert!(inner_projection.is::<ProjectionExec>());
551        assert!(matches!(
552            repartition.partitioning(),
553            Partitioning::Hash(_, 3)
554        ));
555        assert!(repartition.input().is::<CoalescePartitionsExec>());
556        assert_eq!(repartition.input().fetch(), Some(1));
557    }
558
559    fn unordered_input() -> Arc<dyn ExecutionPlan> {
560        let schema = schema();
561        let batch = batch(schema.clone());
562        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
563        Arc::new(TestMemoryExec::try_new(&partitions, schema, None).unwrap())
564    }
565
566    fn ordered_input() -> (Arc<dyn ExecutionPlan>, LexOrdering) {
567        let schema = schema();
568        let ordering = ordering(schema.as_ref(), false);
569        let batch = batch(schema.clone());
570        let partitions = vec![vec![batch.clone()], vec![batch.clone()], vec![batch]];
571        let input = TestMemoryExec::try_new(&partitions, schema, None)
572            .unwrap()
573            .try_with_sort_information(vec![ordering.clone()])
574            .unwrap();
575
576        (Arc::new(input), ordering)
577    }
578
579    fn filter_fetch(input: Arc<dyn ExecutionPlan>, fetch: usize) -> Arc<dyn ExecutionPlan> {
580        Arc::new(
581            FilterExecBuilder::new(lit(true), input)
582                .with_fetch(Some(fetch))
583                .build()
584                .unwrap(),
585        )
586    }
587
588    fn merge_sort_fetch(
589        ordering: LexOrdering,
590        input: Arc<dyn ExecutionPlan>,
591        fetch: usize,
592    ) -> Arc<dyn ExecutionPlan> {
593        Arc::new(MergeSortExec::new(ordering, input, Some(fetch)))
594    }
595
596    fn hash_repartition(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
597        let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3);
598        Arc::new(RepartitionExec::try_new(input, partitioning).unwrap())
599    }
600
601    fn project_a(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
602        Arc::new(
603            ProjectionExec::try_new(
604                vec![(col("a", input.schema().as_ref()).unwrap(), "a".to_string())],
605                input,
606            )
607            .unwrap(),
608        )
609    }
610
611    fn schema() -> Arc<Schema> {
612        Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]))
613    }
614
615    fn batch(schema: Arc<Schema>) -> RecordBatch {
616        RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap()
617    }
618
619    fn ordering(schema: &Schema, descending: bool) -> LexOrdering {
620        LexOrdering::new([PhysicalSortExpr::new(
621            col("a", schema).unwrap(),
622            SortOptions {
623                descending,
624                nulls_first: descending,
625            },
626        )])
627        .unwrap()
628    }
629}
query/optimizer/global_limit.rs

query/optimizer/
global_limit.rs
