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

query/optimizer/
global_limit.rs
