query/dist_plan/
merge_sort.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
15//! Merge sort logical plan for distributed query execution, roughly corresponding to the
16//! `SortPreservingMergeExec` operator in datafusion
17//!
18
19use std::any::Any;
20use std::fmt;
21use std::sync::Arc;
22
23use datafusion::execution::TaskContext;
24use datafusion::physical_plan::execution_plan::CardinalityEffect;
25use datafusion::physical_plan::metrics::MetricsSet;
26use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, update_ordering};
27use datafusion::physical_plan::sorts::sort::SortExec;
28use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
29use datafusion::physical_plan::{
30    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
31    Statistics,
32};
33use datafusion_common::{DataFusionError, Result};
34use datafusion_expr::{Extension, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore};
35use datafusion_physical_expr::{Distribution, LexOrdering, OrderingRequirements};
36
37/// MergeSort Logical Plan, have same field as `Sort`, but indicate it is a merge sort,
38/// which assume each input partition is a sorted stream, and will use `SortPreserveingMergeExec`
39/// to merge them into a single sorted stream.
40#[derive(Hash, PartialOrd, PartialEq, Eq, Clone)]
41pub struct MergeSortLogicalPlan {
42    pub expr: Vec<SortExpr>,
43    pub input: Arc<LogicalPlan>,
44    pub fetch: Option<usize>,
45}
46
47impl fmt::Debug for MergeSortLogicalPlan {
48    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
49        UserDefinedLogicalNodeCore::fmt_for_explain(self, f)
50    }
51}
52
53impl MergeSortLogicalPlan {
54    pub fn new(input: Arc<LogicalPlan>, expr: Vec<SortExpr>, fetch: Option<usize>) -> Self {
55        Self { input, expr, fetch }
56    }
57
58    pub fn name() -> &'static str {
59        "MergeSort"
60    }
61
62    /// Create a [`LogicalPlan::Extension`] node from this merge sort plan
63    pub fn into_logical_plan(self) -> LogicalPlan {
64        LogicalPlan::Extension(Extension {
65            node: Arc::new(self),
66        })
67    }
68}
69
70/// An opaque physical execution node for [`MergeSortLogicalPlan`].
71///
72/// It delegates execution and physical properties to DataFusion's
73/// [`SortPreservingMergeExec`], but intentionally does not expose itself as a
74/// `SortPreservingMergeExec`. `EnforceSorting` is allowed to replace a bare
75/// `SortPreservingMergeExec` with `CoalescePartitionsExec` when the parent does
76/// not require ordering. `MergeSortExec` represents the distributed TopK merge
77/// stage itself, so later physical optimizer rules must not rewrite it into an
78/// unordered fetch.
79#[derive(Debug, Clone)]
80pub(crate) struct MergeSortExec {
81    inner: SortPreservingMergeExec,
82}
83
84impl MergeSortExec {
85    pub(crate) fn new(
86        ordering: LexOrdering,
87        input: Arc<dyn ExecutionPlan>,
88        fetch: Option<usize>,
89    ) -> Self {
90        Self {
91            inner: SortPreservingMergeExec::new(ordering, input).with_fetch(fetch),
92        }
93    }
94
95    fn input_with_fetch(&self, fetch: Option<usize>) -> Arc<dyn ExecutionPlan> {
96        let input = Arc::clone(self.inner.input());
97        if let Some(sort) = input.as_any().downcast_ref::<SortExec>()
98            && sort.preserve_partitioning()
99            && sort.expr() == self.inner.expr()
100        {
101            // Mirror DataFusion's bare SPM plan quality for distributed TopK:
102            // keep the parent `MergeSortExec(fetch)` as the global merge, and
103            // bound the partition-preserving child sort to the same local TopK.
104            // Local top-K is safe because every global top-K row must be within
105            // the top-K rows of its own input partition.
106            Arc::new(sort.with_fetch(fetch))
107        } else {
108            input
109        }
110    }
111}
112
113impl DisplayAs for MergeSortExec {
114    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
115        match t {
116            DisplayFormatType::Default | DisplayFormatType::Verbose => {
117                write!(f, "MergeSortExec: [{}]", self.inner.expr())?;
118                if let Some(fetch) = self.inner.fetch() {
119                    write!(f, ", fetch={fetch}")?;
120                }
121                Ok(())
122            }
123            DisplayFormatType::TreeRender => {
124                if let Some(fetch) = self.inner.fetch() {
125                    writeln!(f, "limit={fetch}")?;
126                }
127
128                for (i, expr) in self.inner.expr().iter().enumerate() {
129                    expr.fmt_sql(f)?;
130                    if i != self.inner.expr().len() - 1 {
131                        write!(f, ", ")?;
132                    }
133                }
134
135                Ok(())
136            }
137        }
138    }
139}
140
141impl ExecutionPlan for MergeSortExec {
142    fn name(&self) -> &str {
143        "MergeSortExec"
144    }
145
146    /// Keeps this node intentionally opaque to DataFusion's type-specialized
147    /// optimizer rewrites.
148    ///
149    /// `MergeSortExec` delegates most behavior to DataFusion's
150    /// `SortPreservingMergeExec`, but it must not expose itself as that type.
151    /// DataFusion's `EnforceSorting` optimizer recognizes a bare
152    /// `SortPreservingMergeExec` via `as_any().downcast_ref::<...>()` and may
153    /// replace it with an unordered `CoalescePartitionsExec(fetch)` when the
154    /// parent does not require sorted output.
155    ///
156    /// That rewrite is valid for an ordinary SPM used only to satisfy parent
157    /// ordering, but not for GreptimeDB's distributed TopK merge stage. In a
158    /// scalar-subquery shape like `ORDER BY ts DESC LIMIT 1`, this node is the
159    /// operator that merges region-local TopK streams into the global TopK.
160    /// Replacing it with unordered coalescing can return a partial/latest row
161    /// from one region instead of the global latest row.
162    ///
163    /// `required_input_ordering()` separately tells DataFusion what ordering this
164    /// node needs from its child, so `EnforceSorting` can insert a `SortExec`
165    /// below `MergeSortExec` when `MergeScanExec` cannot preserve per-partition
166    /// ordering. This opacity is specifically about protecting the merge stage
167    /// itself from the `EnforceSorting` rewrite above.
168    fn as_any(&self) -> &dyn Any {
169        self
170    }
171
172    fn properties(&self) -> &Arc<PlanProperties> {
173        self.inner.properties()
174    }
175
176    /// Forwards DataFusion's order-preserving scan hint through this wrapper.
177    ///
178    /// This mirrors `SortPreservingMergeExec::with_preserve_order()`: if the
179    /// child can produce an order-preserving variant, rebuild the same merge
180    /// stage on top of that child. The returned plan must stay a
181    /// `MergeSortExec`, not a bare SPM, so the distributed TopK merge remains
182    /// opaque to `EnforceSorting`'s SPM-specific rewrite.
183    fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
184        self.inner
185            .input()
186            .with_preserve_order(preserve_order)
187            .map(|new_input| {
188                Arc::new(Self::new(
189                    self.inner.expr().clone(),
190                    new_input,
191                    self.inner.fetch(),
192                )) as Arc<dyn ExecutionPlan>
193            })
194    }
195
196    fn required_input_distribution(&self) -> Vec<Distribution> {
197        self.inner.required_input_distribution()
198    }
199
200    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
201        self.inner.benefits_from_input_partitioning()
202    }
203
204    /// Tells DataFusion that `MergeSortExec` requires each input partition to be
205    /// ordered. This is the contract that makes `EnforceSorting` insert a
206    /// `SortExec` below `MergeSortExec` when the input cannot preserve ordering.
207    ///
208    /// The opacity of `MergeSortExec::as_any`, not this requirement, is what
209    /// prevents DataFusion from rewriting the merge stage itself as a bare
210    /// `SortPreservingMergeExec`.
211    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
212        vec![Some(OrderingRequirements::from(self.inner.expr().clone()))]
213    }
214
215    fn maintains_input_order(&self) -> Vec<bool> {
216        self.inner.maintains_input_order()
217    }
218
219    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
220        self.inner.children()
221    }
222
223    fn with_new_children(
224        self: Arc<Self>,
225        mut children: Vec<Arc<dyn ExecutionPlan>>,
226    ) -> Result<Arc<dyn ExecutionPlan>> {
227        if children.len() != 1 {
228            return Err(DataFusionError::Internal(format!(
229                "MergeSortExec expects exactly one child, got {}",
230                children.len()
231            )));
232        }
233
234        Ok(Arc::new(Self::new(
235            self.inner.expr().clone(),
236            children.swap_remove(0),
237            self.inner.fetch(),
238        )))
239    }
240
241    fn execute(
242        &self,
243        partition: usize,
244        context: Arc<TaskContext>,
245    ) -> Result<SendableRecordBatchStream> {
246        self.inner.execute(partition, context)
247    }
248
249    fn metrics(&self) -> Option<MetricsSet> {
250        self.inner.metrics()
251    }
252
253    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
254        self.inner.partition_statistics(partition)
255    }
256
257    fn cardinality_effect(&self) -> CardinalityEffect {
258        self.inner.cardinality_effect()
259    }
260
261    /// Intentionally keeps DataFusion's generic limit pushdown disabled.
262    ///
263    /// `MergeSortExec` still supports its own global fetch through
264    /// `with_fetch()`. What we must not allow is pushing an external limit below
265    /// this required distributed TopK merge. DataFusion's limit pushdown rules
266    /// know how to treat a bare `SortPreservingMergeExec` as a
267    /// partition-combining node, but `MergeSortExec` is intentionally opaque to
268    /// those SPM-specific downcasts. Enabling generic limit pushdown without also
269    /// teaching the optimizer about this wrapper could return partition-local
270    /// rows instead of the global TopK.
271    fn supports_limit_pushdown(&self) -> bool {
272        false
273    }
274
275    fn fetch(&self) -> Option<usize> {
276        self.inner.fetch()
277    }
278
279    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
280        Some(Arc::new(Self::new(
281            self.inner.expr().clone(),
282            self.input_with_fetch(limit),
283            limit,
284        )))
285    }
286
287    /// Lets DataFusion push a projection below this merge when it can rewrite
288    /// the ordering expressions safely.
289    ///
290    /// This mirrors `SortPreservingMergeExec::try_swapping_with_projection()`
291    /// for plan quality, but re-wraps the result as `MergeSortExec` so the
292    /// distributed merge stage keeps its type identity and opacity.
293    fn try_swapping_with_projection(
294        &self,
295        projection: &ProjectionExec,
296    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
297        if projection.expr().len() >= projection.input().schema().fields().len() {
298            return Ok(None);
299        }
300
301        let Some(updated_exprs) = update_ordering(self.inner.expr().clone(), projection.expr())?
302        else {
303            return Ok(None);
304        };
305
306        Ok(Some(Arc::new(Self::new(
307            updated_exprs,
308            make_with_child(projection, self.inner.input())?,
309            self.inner.fetch(),
310        ))))
311    }
312}
313
314impl UserDefinedLogicalNodeCore for MergeSortLogicalPlan {
315    fn name(&self) -> &str {
316        Self::name()
317    }
318
319    // Allow optimization here
320    fn inputs(&self) -> Vec<&LogicalPlan> {
321        vec![self.input.as_ref()]
322    }
323
324    fn schema(&self) -> &datafusion_common::DFSchemaRef {
325        self.input.schema()
326    }
327
328    // Allow further optimization
329    fn expressions(&self) -> Vec<datafusion_expr::Expr> {
330        self.expr.iter().map(|sort| sort.expr.clone()).collect()
331    }
332
333    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
334        write!(f, "MergeSort: ")?;
335        for (i, expr_item) in self.expr.iter().enumerate() {
336            if i > 0 {
337                write!(f, ", ")?;
338            }
339            write!(f, "{expr_item}")?;
340        }
341        if let Some(a) = self.fetch {
342            write!(f, ", fetch={a}")?;
343        }
344        Ok(())
345    }
346
347    fn with_exprs_and_inputs(
348        &self,
349        exprs: Vec<datafusion::prelude::Expr>,
350        mut inputs: Vec<LogicalPlan>,
351    ) -> Result<Self> {
352        let mut zelf = self.clone();
353        zelf.expr = zelf
354            .expr
355            .into_iter()
356            .zip(exprs)
357            .map(|(sort, expr)| sort.with_expr(expr))
358            .collect();
359        zelf.input = Arc::new(inputs.pop().ok_or_else(|| {
360            DataFusionError::Internal("Expected exactly one input with MergeSort".to_string())
361        })?);
362        Ok(zelf)
363    }
364}
365
366/// Turn `Sort` into `MergeSort` if possible
367pub fn merge_sort_transformer(plan: &LogicalPlan) -> Option<LogicalPlan> {
368    if let LogicalPlan::Sort(sort) = plan {
369        Some(
370            MergeSortLogicalPlan::new(sort.input.clone(), sort.expr.clone(), sort.fetch)
371                .into_logical_plan(),
372        )
373    } else {
374        None
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use arrow_schema::{DataType, Field, Schema, SortOptions};
381    use datafusion::physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{
382        OrderPreservationContext, plan_with_order_breaking_variants,
383    };
384    use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
385    use datafusion::physical_plan::displayable;
386    use datafusion::physical_plan::empty::EmptyExec;
387    use datafusion_physical_expr::PhysicalSortExpr;
388    use datafusion_physical_expr::expressions::col as physical_col;
389
390    use super::*;
391
392    /// Test double that records DataFusion's preserve-order signal while
393    /// otherwise behaving like a transparent wrapper around its child.
394    #[derive(Debug, Clone)]
395    struct PreserveOrderProbeExec {
396        inner: Arc<dyn ExecutionPlan>,
397        preserve_order: bool,
398    }
399
400    impl PreserveOrderProbeExec {
401        fn new(inner: Arc<dyn ExecutionPlan>) -> Self {
402            Self {
403                inner,
404                preserve_order: false,
405            }
406        }
407    }
408
409    impl DisplayAs for PreserveOrderProbeExec {
410        fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
411            write!(
412                f,
413                "PreserveOrderProbeExec: preserve_order={}",
414                self.preserve_order
415            )
416        }
417    }
418
419    impl ExecutionPlan for PreserveOrderProbeExec {
420        fn name(&self) -> &str {
421            "PreserveOrderProbeExec"
422        }
423
424        fn as_any(&self) -> &dyn Any {
425            self
426        }
427
428        fn properties(&self) -> &Arc<PlanProperties> {
429            self.inner.properties()
430        }
431
432        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
433            vec![&self.inner]
434        }
435
436        fn with_new_children(
437            self: Arc<Self>,
438            mut children: Vec<Arc<dyn ExecutionPlan>>,
439        ) -> Result<Arc<dyn ExecutionPlan>> {
440            if children.len() != 1 {
441                return Err(DataFusionError::Internal(format!(
442                    "PreserveOrderProbeExec expects exactly one child, got {}",
443                    children.len()
444                )));
445            }
446
447            Ok(Arc::new(Self {
448                inner: children.swap_remove(0),
449                preserve_order: self.preserve_order,
450            }))
451        }
452
453        fn execute(
454            &self,
455            partition: usize,
456            context: Arc<TaskContext>,
457        ) -> Result<SendableRecordBatchStream> {
458            self.inner.execute(partition, context)
459        }
460
461        fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
462            Some(Arc::new(Self {
463                inner: Arc::clone(&self.inner),
464                preserve_order,
465            }))
466        }
467    }
468
469    fn test_ordering(schema: &Schema) -> LexOrdering {
470        LexOrdering::new([PhysicalSortExpr::new(
471            physical_col("ts", schema).unwrap(),
472            SortOptions {
473                descending: true,
474                nulls_first: false,
475            },
476        )])
477        .unwrap()
478    }
479
480    #[test]
481    fn merge_sort_exec_is_opaque_and_preserves_topk_requirements() {
482        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
483        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
484        let ordering = test_ordering(schema.as_ref());
485
486        let merge_sort =
487            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
488
489        assert_eq!(merge_sort.name(), "MergeSortExec");
490        assert!(
491            merge_sort
492                .as_any()
493                .downcast_ref::<SortPreservingMergeExec>()
494                .is_none(),
495            "MergeSortExec must stay opaque to EnforceSorting's bare SortPreservingMerge rewrite"
496        );
497        assert_eq!(merge_sort.fetch(), Some(1));
498        assert!(!merge_sort.supports_limit_pushdown());
499        assert!(merge_sort.required_input_ordering()[0].is_some());
500
501        let tree = displayable(merge_sort.as_ref()).tree_render().to_string();
502        assert!(tree.contains("MergeSortExec"));
503        assert!(!tree.contains("SortPreservingMergeExec"));
504
505        let fetched = merge_sort.with_fetch(Some(2)).unwrap();
506        assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
507        assert_eq!(fetched.fetch(), Some(2));
508    }
509
510    #[test]
511    fn merge_sort_exec_required_input_ordering_matches_spm() {
512        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
513        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
514        let ordering = test_ordering(schema.as_ref());
515
516        let merge_sort = MergeSortExec::new(ordering.clone(), Arc::clone(&input), Some(1));
517        let bare_spm =
518            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1));
519
520        assert_eq!(
521            merge_sort.required_input_ordering(),
522            vec![Some(OrderingRequirements::from(ordering))],
523            "MergeSortExec must require locally sorted input partitions for the merge key"
524        );
525        assert_eq!(
526            merge_sort.required_input_ordering(),
527            bare_spm.required_input_ordering(),
528            "MergeSortExec's child ordering contract should mirror SortPreservingMergeExec"
529        );
530        assert_eq!(
531            merge_sort.maintains_input_order(),
532            bare_spm.maintains_input_order()
533        );
534    }
535
536    #[test]
537    fn merge_sort_exec_with_fetch_pushes_fetch_to_child_sort() {
538        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
539        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
540        let ordering = test_ordering(schema.as_ref());
541        let child_sort =
542            Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true))
543                as Arc<dyn ExecutionPlan>;
544        let merge_sort = MergeSortExec::new(ordering, child_sort, None);
545
546        let fetched = merge_sort.with_fetch(Some(2)).unwrap();
547
548        assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
549        assert_eq!(fetched.fetch(), Some(2));
550        let child_sort = fetched.children()[0]
551            .as_any()
552            .downcast_ref::<SortExec>()
553            .unwrap();
554        assert_eq!(child_sort.fetch(), Some(2));
555        assert!(child_sort.preserve_partitioning());
556    }
557
558    #[test]
559    fn merge_sort_exec_with_preserve_order_matches_spm_but_keeps_wrapper() {
560        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
561        let input = Arc::new(PreserveOrderProbeExec::new(Arc::new(
562            EmptyExec::new(schema.clone()).with_partitions(2),
563        ))) as _;
564        let ordering = test_ordering(schema.as_ref());
565
566        let bare_spm =
567            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1));
568        let preserved_spm = bare_spm.with_preserve_order(true).unwrap();
569        assert!(
570            preserved_spm
571                .as_any()
572                .downcast_ref::<SortPreservingMergeExec>()
573                .is_some(),
574            "bare SPM should rebuild as bare SPM"
575        );
576        assert!(
577            preserved_spm.children()[0]
578                .as_any()
579                .downcast_ref::<PreserveOrderProbeExec>()
580                .unwrap()
581                .preserve_order
582        );
583
584        let merge_sort = MergeSortExec::new(ordering, input, Some(1));
585        let preserved_merge_sort = merge_sort.with_preserve_order(true).unwrap();
586        assert!(
587            preserved_merge_sort
588                .as_any()
589                .downcast_ref::<MergeSortExec>()
590                .is_some(),
591            "MergeSortExec must rewrap the preserve-order child as MergeSortExec"
592        );
593        assert!(
594            preserved_merge_sort
595                .as_any()
596                .downcast_ref::<SortPreservingMergeExec>()
597                .is_none(),
598            "MergeSortExec must not expose a bare SPM after with_preserve_order"
599        );
600        assert_eq!(preserved_merge_sort.fetch(), Some(1));
601        assert_eq!(
602            preserved_merge_sort.required_input_ordering(),
603            preserved_spm.required_input_ordering(),
604            "preserve-order rewrite should keep the same SPM child-ordering contract"
605        );
606        assert!(
607            preserved_merge_sort.children()[0]
608                .as_any()
609                .downcast_ref::<PreserveOrderProbeExec>()
610                .unwrap()
611                .preserve_order
612        );
613    }
614
615    #[test]
616    fn merge_sort_exec_projection_swap_matches_spm_but_keeps_wrapper() -> Result<()> {
617        let schema = Arc::new(Schema::new(vec![
618            Field::new("value", DataType::Int64, false),
619            Field::new("ts", DataType::Int64, false),
620            Field::new("tag", DataType::Utf8, false),
621        ]));
622        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
623        let ordering = test_ordering(schema.as_ref());
624
625        let bare_spm = Arc::new(
626            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)),
627        ) as Arc<dyn ExecutionPlan>;
628        let spm_projection = ProjectionExec::try_new(
629            vec![
630                (physical_col("ts", schema.as_ref())?, "ts".to_string()),
631                (physical_col("tag", schema.as_ref())?, "tag".to_string()),
632            ],
633            Arc::clone(&bare_spm),
634        )?;
635        let swapped_spm = bare_spm
636            .try_swapping_with_projection(&spm_projection)?
637            .expect("SPM should accept a narrowing projection that preserves the sort key");
638        assert!(
639            swapped_spm
640                .as_any()
641                .downcast_ref::<SortPreservingMergeExec>()
642                .is_some(),
643            "bare SPM should rebuild as bare SPM"
644        );
645
646        let merge_sort =
647            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
648        let merge_projection = ProjectionExec::try_new(
649            vec![
650                (physical_col("ts", schema.as_ref())?, "ts".to_string()),
651                (physical_col("tag", schema.as_ref())?, "tag".to_string()),
652            ],
653            Arc::clone(&merge_sort),
654        )?;
655        let swapped_merge_sort = merge_sort
656            .try_swapping_with_projection(&merge_projection)?
657            .expect("MergeSortExec should accept the same projection swap as SPM");
658
659        assert!(
660            swapped_merge_sort
661                .as_any()
662                .downcast_ref::<MergeSortExec>()
663                .is_some(),
664            "MergeSortExec must rewrap projection swaps as MergeSortExec"
665        );
666        assert!(
667            swapped_merge_sort
668                .as_any()
669                .downcast_ref::<SortPreservingMergeExec>()
670                .is_none(),
671            "MergeSortExec must not expose a bare SPM after projection swap"
672        );
673        assert_eq!(swapped_merge_sort.fetch(), Some(1));
674        assert!(
675            swapped_merge_sort.children()[0]
676                .as_any()
677                .downcast_ref::<ProjectionExec>()
678                .is_some(),
679            "the projection should move below MergeSortExec"
680        );
681        let swapped_schema = swapped_merge_sort.schema();
682        assert_eq!(
683            swapped_schema
684                .fields()
685                .iter()
686                .map(|field| field.name().as_str())
687                .collect::<Vec<_>>(),
688            vec!["ts", "tag"],
689            "swapped MergeSortExec should expose the projected schema"
690        );
691
692        let projected_ordering = LexOrdering::new([PhysicalSortExpr::new(
693            physical_col("ts", swapped_merge_sort.children()[0].schema().as_ref())?,
694            SortOptions {
695                descending: true,
696                nulls_first: false,
697            },
698        )])
699        .unwrap();
700        assert_eq!(
701            swapped_merge_sort.required_input_ordering(),
702            vec![Some(OrderingRequirements::from(projected_ordering))],
703            "projection swap must rewrite the ordering to the child projection's schema"
704        );
705        assert_eq!(
706            swapped_merge_sort.required_input_ordering(),
707            swapped_spm.required_input_ordering(),
708            "MergeSortExec projection swap should mirror SPM's ordering rewrite"
709        );
710
711        let spm_projection_without_sort_key = ProjectionExec::try_new(
712            vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())],
713            Arc::clone(&bare_spm),
714        )?;
715        let merge_projection_without_sort_key = ProjectionExec::try_new(
716            vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())],
717            Arc::clone(&merge_sort),
718        )?;
719        assert!(
720            bare_spm
721                .try_swapping_with_projection(&spm_projection_without_sort_key)?
722                .is_none(),
723            "SPM must reject projection swaps that drop the sort key"
724        );
725        assert!(
726            merge_sort
727                .try_swapping_with_projection(&merge_projection_without_sort_key)?
728                .is_none(),
729            "MergeSortExec should reject the same projection swap as SPM"
730        );
731
732        Ok(())
733    }
734
735    #[test]
736    fn enforce_sorting_rewrite_keeps_merge_sort_exec_opaque() {
737        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
738        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
739        let ordering = test_ordering(schema.as_ref());
740
741        let bare_spm = Arc::new(
742            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)),
743        ) as Arc<dyn ExecutionPlan>;
744        let optimized_spm = plan_with_order_breaking_variants(OrderPreservationContext::new(
745            bare_spm,
746            false,
747            vec![OrderPreservationContext::new(
748                Arc::clone(&input),
749                false,
750                vec![],
751            )],
752        ))
753        .unwrap()
754        .plan;
755        assert!(
756            optimized_spm
757                .as_any()
758                .downcast_ref::<CoalescePartitionsExec>()
759                .is_some(),
760            "this regression test must exercise EnforceSorting's bare SPM -> CoalescePartitionsExec rewrite"
761        );
762
763        let merge_sort =
764            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
765        let optimized_merge_sort =
766            plan_with_order_breaking_variants(OrderPreservationContext::new(
767                Arc::clone(&merge_sort),
768                false,
769                vec![OrderPreservationContext::new(
770                    Arc::clone(merge_sort.children()[0]),
771                    false,
772                    vec![],
773                )],
774            ))
775            .unwrap()
776            .plan;
777        assert!(
778            optimized_merge_sort
779                .as_any()
780                .downcast_ref::<MergeSortExec>()
781                .is_some(),
782            "MergeSortExec must stay opaque to the bare SPM rewrite"
783        );
784        assert!(
785            optimized_merge_sort
786                .as_any()
787                .downcast_ref::<CoalescePartitionsExec>()
788                .is_none(),
789            "MergeSortExec(fetch) is the required distributed TopK merge stage, not an unordered coalesce"
790        );
791        assert_eq!(optimized_merge_sort.fetch(), Some(1));
792    }
793}
query/dist_plan/merge_sort.rs

query/dist_plan/
merge_sort.rs
