fix: preserve distributed topk merge ordering (#8432)

* fix: preserve distributed topk for scalar latest Signed-off-by: discord9 <discord9@163.com> * fix: preserve scalar latest ordering across merge scan Signed-off-by: discord9 <discord9@163.com> * fix: carry merge scan ordering from planner rewrite Signed-off-by: discord9 <discord9@163.com> * test: cover merge scan ordering metadata Signed-off-by: discord9 <discord9@163.com> * test: cover merge scan partition ordering gate Signed-off-by: discord9 <discord9@163.com> * fix: require adjacent merge sort for ordering metadata Signed-off-by: discord9 <discord9@163.com> * test: update ordering sqlness plans Signed-off-by: discord9 <discord9@163.com> * fix: preserve distributed topk merge ordering Signed-off-by: discord9 <discord9@163.com> * test: cover latest per series queries Signed-off-by: discord9 <discord9@163.com> * test: update distributed merge sort sqlness Signed-off-by: discord9 <discord9@163.com> * test: cover merge scan ordering partition gate Signed-off-by: discord9 <discord9@163.com> * test: cover merge scan over-partition ordering Signed-off-by: discord9 <discord9@163.com> * test: cover distributed latest with low parallelism Signed-off-by: discord9 <discord9@163.com> * fix: keep distributed merge sort opaque Signed-off-by: discord9 <discord9@163.com> * test: assert distributed scalar latest merge sort Signed-off-by: discord9 <discord9@163.com> * test: clarify merge scan ordering helper Signed-off-by: discord9 <discord9@163.com> * fix: complete merge sort exec delegation Signed-off-by: discord9 <discord9@163.com> * docs: explain merge sort limit pushdown Signed-off-by: discord9 <discord9@163.com> * test: cover merge sort optimizer opacity Signed-off-by: discord9 <discord9@163.com> * docs: explain merge sort optimizer hooks Signed-off-by: discord9 <discord9@163.com> * test: cover merge sort optimizer hooks Signed-off-by: discord9 <discord9@163.com> * fix: preserve merge sort child topk Signed-off-by: discord9 <discord9@163.com> * test: update order by topk plan Signed-off-by: discord9 <discord9@163.com> * fix: recognize merge sort global fetch Signed-off-by: discord9 <discord9@163.com> --------- Signed-off-by: discord9 <discord9@163.com>
2026-08-18 03:58:29 +00:00 · 2026-07-14 20:53:07 +08:00
parent b513dbaf4a
commit b30765e5e9
21 changed files with 1447 additions and 73 deletions
@@ -27,6 +27,7 @@ mod remote_dyn_filter_registry;
 pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions};
 pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId};
 pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan};
+pub(crate) use merge_sort::MergeSortExec;
 pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner};
 pub use predicate_extractor::PredicateExtractor;
 pub use region_pruner::ConstraintPruner;
@@ -1031,7 +1031,7 @@ mod tests {
    use datafusion::execution::SessionStateBuilder;
    use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
    use datafusion_common::TableReference;
-    use datafusion_expr::{LogicalPlanBuilder, lit};
+    use datafusion_expr::{LogicalPlanBuilder, col, lit};
    use datafusion_physical_expr::Distribution;
    use datafusion_physical_expr::expressions::{
        Column, DynamicFilterPhysicalExpr, lit as physical_lit,
@@ -1051,6 +1051,65 @@ mod tests {
        QueryId::from(Uuid::from_u128(value))
    }

+    fn merge_scan_exec_with_sorted_input(
+        region_count: u64,
+        target_partition: usize,
+    ) -> MergeScanExec {
+        let session_state = SessionStateBuilder::new().build();
+        let plan = LogicalPlanBuilder::empty(true)
+            .project(vec![lit(1i64).alias("ts")])
+            .unwrap()
+            .sort(vec![col("ts").sort(false, true)])
+            .unwrap()
+            .build()
+            .unwrap();
+        let schema = plan.schema().as_arrow().clone();
+        let regions = (0..region_count)
+            .map(|region_number| RegionId::new(1024, region_number as u32))
+            .collect();
+
+        MergeScanExec::new(
+            &session_state,
+            // The table name is not relevant to these ordering metadata tests;
+            // `MergeScanExec::new` requires one to model the production plan.
+            TableName::new("catalog", "schema", "table"),
+            regions,
+            plan,
+            &schema,
+            Arc::new(TestRegionQueryHandler),
+            QueryContext::arc(),
+            target_partition,
+            AliasMapping::new(),
+            None,
+            false,
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn merge_scan_does_not_advertise_ordering_when_partition_may_merge_regions() {
+        let exec = merge_scan_exec_with_sorted_input(3, 2);
+
+        assert!(
+            exec.properties().output_ordering().is_none(),
+            "target_partition < region_count means one output partition may concatenate multiple sorted region streams"
+        );
+    }
+
+    #[test]
+    fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() {
+        let exec = merge_scan_exec_with_sorted_input(3, 3);
+
+        assert!(exec.properties().output_ordering().is_some());
+    }
+
+    #[test]
+    fn merge_scan_advertises_ordering_when_partitions_exceed_regions() {
+        let exec = merge_scan_exec_with_sorted_input(3, 4);
+
+        assert!(exec.properties().output_ordering().is_some());
+    }
+
    #[test]
    fn remote_dyn_filter_region_query_context_registers_before_do_get() {
        let registry_manager = Arc::new(DynFilterRegistryManager::default());
@@ -16,11 +16,22 @@
 //! `SortPreservingMergeExec` operator in datafusion
 //!

+use std::any::Any;
 use std::fmt;
 use std::sync::Arc;

+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::metrics::MetricsSet;
+use datafusion::physical_plan::projection::{ProjectionExec, make_with_child, update_ordering};
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
+    Statistics,
+};
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::{Extension, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore};
+use datafusion_physical_expr::{Distribution, LexOrdering, OrderingRequirements};

 /// MergeSort Logical Plan, have same field as `Sort`, but indicate it is a merge sort,
 /// which assume each input partition is a sorted stream, and will use `SortPreserveingMergeExec`
@@ -53,14 +64,245 @@ impl MergeSortLogicalPlan {
            node: Arc::new(self),
        })
    }
+}

-    /// Convert self to a [`Sort`] logical plan with same input and expressions
-    pub fn into_sort(self) -> LogicalPlan {
-        LogicalPlan::Sort(datafusion::logical_expr::Sort {
-            input: self.input.clone(),
-            expr: self.expr,
-            fetch: self.fetch,
-        })
+/// An opaque physical execution node for [`MergeSortLogicalPlan`].
+///
+/// It delegates execution and physical properties to DataFusion's
+/// [`SortPreservingMergeExec`], but intentionally does not expose itself as a
+/// `SortPreservingMergeExec`. `EnforceSorting` is allowed to replace a bare
+/// `SortPreservingMergeExec` with `CoalescePartitionsExec` when the parent does
+/// not require ordering. `MergeSortExec` represents the distributed TopK merge
+/// stage itself, so later physical optimizer rules must not rewrite it into an
+/// unordered fetch.
+#[derive(Debug, Clone)]
+pub(crate) struct MergeSortExec {
+    inner: SortPreservingMergeExec,
+}
+
+impl MergeSortExec {
+    pub(crate) fn new(
+        ordering: LexOrdering,
+        input: Arc<dyn ExecutionPlan>,
+        fetch: Option<usize>,
+    ) -> Self {
+        Self {
+            inner: SortPreservingMergeExec::new(ordering, input).with_fetch(fetch),
+        }
+    }
+
+    fn input_with_fetch(&self, fetch: Option<usize>) -> Arc<dyn ExecutionPlan> {
+        let input = Arc::clone(self.inner.input());
+        if let Some(sort) = input.as_any().downcast_ref::<SortExec>()
+            && sort.preserve_partitioning()
+            && sort.expr() == self.inner.expr()
+        {
+            // Mirror DataFusion's bare SPM plan quality for distributed TopK:
+            // keep the parent `MergeSortExec(fetch)` as the global merge, and
+            // bound the partition-preserving child sort to the same local TopK.
+            // Local top-K is safe because every global top-K row must be within
+            // the top-K rows of its own input partition.
+            Arc::new(sort.with_fetch(fetch))
+        } else {
+            input
+        }
+    }
+}
+
+impl DisplayAs for MergeSortExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "MergeSortExec: [{}]", self.inner.expr())?;
+                if let Some(fetch) = self.inner.fetch() {
+                    write!(f, ", fetch={fetch}")?;
+                }
+                Ok(())
+            }
+            DisplayFormatType::TreeRender => {
+                if let Some(fetch) = self.inner.fetch() {
+                    writeln!(f, "limit={fetch}")?;
+                }
+
+                for (i, expr) in self.inner.expr().iter().enumerate() {
+                    expr.fmt_sql(f)?;
+                    if i != self.inner.expr().len() - 1 {
+                        write!(f, ", ")?;
+                    }
+                }
+
+                Ok(())
+            }
+        }
+    }
+}
+
+impl ExecutionPlan for MergeSortExec {
+    fn name(&self) -> &str {
+        "MergeSortExec"
+    }
+
+    /// Keeps this node intentionally opaque to DataFusion's type-specialized
+    /// optimizer rewrites.
+    ///
+    /// `MergeSortExec` delegates most behavior to DataFusion's
+    /// `SortPreservingMergeExec`, but it must not expose itself as that type.
+    /// DataFusion's `EnforceSorting` optimizer recognizes a bare
+    /// `SortPreservingMergeExec` via `as_any().downcast_ref::<...>()` and may
+    /// replace it with an unordered `CoalescePartitionsExec(fetch)` when the
+    /// parent does not require sorted output.
+    ///
+    /// That rewrite is valid for an ordinary SPM used only to satisfy parent
+    /// ordering, but not for GreptimeDB's distributed TopK merge stage. In a
+    /// scalar-subquery shape like `ORDER BY ts DESC LIMIT 1`, this node is the
+    /// operator that merges region-local TopK streams into the global TopK.
+    /// Replacing it with unordered coalescing can return a partial/latest row
+    /// from one region instead of the global latest row.
+    ///
+    /// `required_input_ordering()` separately tells DataFusion what ordering this
+    /// node needs from its child, so `EnforceSorting` can insert a `SortExec`
+    /// below `MergeSortExec` when `MergeScanExec` cannot preserve per-partition
+    /// ordering. This opacity is specifically about protecting the merge stage
+    /// itself from the `EnforceSorting` rewrite above.
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        self.inner.properties()
+    }
+
+    /// Forwards DataFusion's order-preserving scan hint through this wrapper.
+    ///
+    /// This mirrors `SortPreservingMergeExec::with_preserve_order()`: if the
+    /// child can produce an order-preserving variant, rebuild the same merge
+    /// stage on top of that child. The returned plan must stay a
+    /// `MergeSortExec`, not a bare SPM, so the distributed TopK merge remains
+    /// opaque to `EnforceSorting`'s SPM-specific rewrite.
+    fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
+        self.inner
+            .input()
+            .with_preserve_order(preserve_order)
+            .map(|new_input| {
+                Arc::new(Self::new(
+                    self.inner.expr().clone(),
+                    new_input,
+                    self.inner.fetch(),
+                )) as Arc<dyn ExecutionPlan>
+            })
+    }
+
+    fn required_input_distribution(&self) -> Vec<Distribution> {
+        self.inner.required_input_distribution()
+    }
+
+    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+        self.inner.benefits_from_input_partitioning()
+    }
+
+    /// Tells DataFusion that `MergeSortExec` requires each input partition to be
+    /// ordered. This is the contract that makes `EnforceSorting` insert a
+    /// `SortExec` below `MergeSortExec` when the input cannot preserve ordering.
+    ///
+    /// The opacity of `MergeSortExec::as_any`, not this requirement, is what
+    /// prevents DataFusion from rewriting the merge stage itself as a bare
+    /// `SortPreservingMergeExec`.
+    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+        vec![Some(OrderingRequirements::from(self.inner.expr().clone()))]
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.inner.maintains_input_order()
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        self.inner.children()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        mut children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if children.len() != 1 {
+            return Err(DataFusionError::Internal(format!(
+                "MergeSortExec expects exactly one child, got {}",
+                children.len()
+            )));
+        }
+
+        Ok(Arc::new(Self::new(
+            self.inner.expr().clone(),
+            children.swap_remove(0),
+            self.inner.fetch(),
+        )))
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        self.inner.execute(partition, context)
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        self.inner.metrics()
+    }
+
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.inner.partition_statistics(partition)
+    }
+
+    /// Intentionally keeps DataFusion's generic limit pushdown disabled.
+    ///
+    /// `MergeSortExec` still supports its own global fetch through
+    /// `with_fetch()`. What we must not allow is pushing an external limit below
+    /// this required distributed TopK merge. DataFusion's limit pushdown rules
+    /// know how to treat a bare `SortPreservingMergeExec` as a
+    /// partition-combining node, but `MergeSortExec` is intentionally opaque to
+    /// those SPM-specific downcasts. Enabling generic limit pushdown without also
+    /// teaching the optimizer about this wrapper could return partition-local
+    /// rows instead of the global TopK.
+    fn supports_limit_pushdown(&self) -> bool {
+        false
+    }
+
+    fn fetch(&self) -> Option<usize> {
+        self.inner.fetch()
+    }
+
+    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+        Some(Arc::new(Self::new(
+            self.inner.expr().clone(),
+            self.input_with_fetch(limit),
+            limit,
+        )))
+    }
+
+    /// Lets DataFusion push a projection below this merge when it can rewrite
+    /// the ordering expressions safely.
+    ///
+    /// This mirrors `SortPreservingMergeExec::try_swapping_with_projection()`
+    /// for plan quality, but re-wraps the result as `MergeSortExec` so the
+    /// distributed merge stage keeps its type identity and opacity.
+    fn try_swapping_with_projection(
+        &self,
+        projection: &ProjectionExec,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if projection.expr().len() >= projection.input().schema().fields().len() {
+            return Ok(None);
+        }
+
+        let Some(updated_exprs) = update_ordering(self.inner.expr().clone(), projection.expr())?
+        else {
+            return Ok(None);
+        };
+
+        Ok(Some(Arc::new(Self::new(
+            updated_exprs,
+            make_with_child(projection, self.inner.input())?,
+            self.inner.fetch(),
+        ))))
    }
 }

@@ -127,3 +369,420 @@ pub fn merge_sort_transformer(plan: &LogicalPlan) -> Option<LogicalPlan> {
        None
    }
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow_schema::{DataType, Field, Schema, SortOptions};
+    use datafusion::physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{
+        OrderPreservationContext, plan_with_order_breaking_variants,
+    };
+    use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+    use datafusion::physical_plan::displayable;
+    use datafusion::physical_plan::empty::EmptyExec;
+    use datafusion_physical_expr::PhysicalSortExpr;
+    use datafusion_physical_expr::expressions::col as physical_col;
+
+    use super::*;
+
+    /// Test double that records DataFusion's preserve-order signal while
+    /// otherwise behaving like a transparent wrapper around its child.
+    #[derive(Debug, Clone)]
+    struct PreserveOrderProbeExec {
+        inner: Arc<dyn ExecutionPlan>,
+        preserve_order: bool,
+    }
+
+    impl PreserveOrderProbeExec {
+        fn new(inner: Arc<dyn ExecutionPlan>) -> Self {
+            Self {
+                inner,
+                preserve_order: false,
+            }
+        }
+    }
+
+    impl DisplayAs for PreserveOrderProbeExec {
+        fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(
+                f,
+                "PreserveOrderProbeExec: preserve_order={}",
+                self.preserve_order
+            )
+        }
+    }
+
+    impl ExecutionPlan for PreserveOrderProbeExec {
+        fn name(&self) -> &str {
+            "PreserveOrderProbeExec"
+        }
+
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        fn properties(&self) -> &Arc<PlanProperties> {
+            self.inner.properties()
+        }
+
+        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+            vec![&self.inner]
+        }
+
+        fn with_new_children(
+            self: Arc<Self>,
+            mut children: Vec<Arc<dyn ExecutionPlan>>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            if children.len() != 1 {
+                return Err(DataFusionError::Internal(format!(
+                    "PreserveOrderProbeExec expects exactly one child, got {}",
+                    children.len()
+                )));
+            }
+
+            Ok(Arc::new(Self {
+                inner: children.swap_remove(0),
+                preserve_order: self.preserve_order,
+            }))
+        }
+
+        fn execute(
+            &self,
+            partition: usize,
+            context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            self.inner.execute(partition, context)
+        }
+
+        fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
+            Some(Arc::new(Self {
+                inner: Arc::clone(&self.inner),
+                preserve_order,
+            }))
+        }
+    }
+
+    fn test_ordering(schema: &Schema) -> LexOrdering {
+        LexOrdering::new([PhysicalSortExpr::new(
+            physical_col("ts", schema).unwrap(),
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )])
+        .unwrap()
+    }
+
+    #[test]
+    fn merge_sort_exec_is_opaque_and_preserves_topk_requirements() {
+        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
+        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
+        let ordering = test_ordering(schema.as_ref());
+
+        let merge_sort =
+            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
+
+        assert_eq!(merge_sort.name(), "MergeSortExec");
+        assert!(
+            merge_sort
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_none(),
+            "MergeSortExec must stay opaque to EnforceSorting's bare SortPreservingMerge rewrite"
+        );
+        assert_eq!(merge_sort.fetch(), Some(1));
+        assert!(!merge_sort.supports_limit_pushdown());
+        assert!(merge_sort.required_input_ordering()[0].is_some());
+
+        let tree = displayable(merge_sort.as_ref()).tree_render().to_string();
+        assert!(tree.contains("MergeSortExec"));
+        assert!(!tree.contains("SortPreservingMergeExec"));
+
+        let fetched = merge_sort.with_fetch(Some(2)).unwrap();
+        assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
+        assert_eq!(fetched.fetch(), Some(2));
+    }
+
+    #[test]
+    fn merge_sort_exec_required_input_ordering_matches_spm() {
+        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
+        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
+        let ordering = test_ordering(schema.as_ref());
+
+        let merge_sort = MergeSortExec::new(ordering.clone(), Arc::clone(&input), Some(1));
+        let bare_spm =
+            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1));
+
+        assert_eq!(
+            merge_sort.required_input_ordering(),
+            vec![Some(OrderingRequirements::from(ordering))],
+            "MergeSortExec must require locally sorted input partitions for the merge key"
+        );
+        assert_eq!(
+            merge_sort.required_input_ordering(),
+            bare_spm.required_input_ordering(),
+            "MergeSortExec's child ordering contract should mirror SortPreservingMergeExec"
+        );
+        assert_eq!(
+            merge_sort.maintains_input_order(),
+            bare_spm.maintains_input_order()
+        );
+    }
+
+    #[test]
+    fn merge_sort_exec_with_fetch_pushes_fetch_to_child_sort() {
+        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
+        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
+        let ordering = test_ordering(schema.as_ref());
+        let child_sort =
+            Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true))
+                as Arc<dyn ExecutionPlan>;
+        let merge_sort = MergeSortExec::new(ordering, child_sort, None);
+
+        let fetched = merge_sort.with_fetch(Some(2)).unwrap();
+
+        assert!(fetched.as_any().downcast_ref::<MergeSortExec>().is_some());
+        assert_eq!(fetched.fetch(), Some(2));
+        let child_sort = fetched.children()[0]
+            .as_any()
+            .downcast_ref::<SortExec>()
+            .unwrap();
+        assert_eq!(child_sort.fetch(), Some(2));
+        assert!(child_sort.preserve_partitioning());
+    }
+
+    #[test]
+    fn merge_sort_exec_with_preserve_order_matches_spm_but_keeps_wrapper() {
+        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
+        let input = Arc::new(PreserveOrderProbeExec::new(Arc::new(
+            EmptyExec::new(schema.clone()).with_partitions(2),
+        ))) as _;
+        let ordering = test_ordering(schema.as_ref());
+
+        let bare_spm =
+            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1));
+        let preserved_spm = bare_spm.with_preserve_order(true).unwrap();
+        assert!(
+            preserved_spm
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_some(),
+            "bare SPM should rebuild as bare SPM"
+        );
+        assert!(
+            preserved_spm.children()[0]
+                .as_any()
+                .downcast_ref::<PreserveOrderProbeExec>()
+                .unwrap()
+                .preserve_order
+        );
+
+        let merge_sort = MergeSortExec::new(ordering, input, Some(1));
+        let preserved_merge_sort = merge_sort.with_preserve_order(true).unwrap();
+        assert!(
+            preserved_merge_sort
+                .as_any()
+                .downcast_ref::<MergeSortExec>()
+                .is_some(),
+            "MergeSortExec must rewrap the preserve-order child as MergeSortExec"
+        );
+        assert!(
+            preserved_merge_sort
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_none(),
+            "MergeSortExec must not expose a bare SPM after with_preserve_order"
+        );
+        assert_eq!(preserved_merge_sort.fetch(), Some(1));
+        assert_eq!(
+            preserved_merge_sort.required_input_ordering(),
+            preserved_spm.required_input_ordering(),
+            "preserve-order rewrite should keep the same SPM child-ordering contract"
+        );
+        assert!(
+            preserved_merge_sort.children()[0]
+                .as_any()
+                .downcast_ref::<PreserveOrderProbeExec>()
+                .unwrap()
+                .preserve_order
+        );
+    }
+
+    #[test]
+    fn merge_sort_exec_projection_swap_matches_spm_but_keeps_wrapper() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("value", DataType::Int64, false),
+            Field::new("ts", DataType::Int64, false),
+            Field::new("tag", DataType::Utf8, false),
+        ]));
+        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
+        let ordering = test_ordering(schema.as_ref());
+
+        let bare_spm = Arc::new(
+            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)),
+        ) as Arc<dyn ExecutionPlan>;
+        let spm_projection = ProjectionExec::try_new(
+            vec![
+                (physical_col("ts", schema.as_ref())?, "ts".to_string()),
+                (physical_col("tag", schema.as_ref())?, "tag".to_string()),
+            ],
+            Arc::clone(&bare_spm),
+        )?;
+        let swapped_spm = bare_spm
+            .try_swapping_with_projection(&spm_projection)?
+            .expect("SPM should accept a narrowing projection that preserves the sort key");
+        assert!(
+            swapped_spm
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_some(),
+            "bare SPM should rebuild as bare SPM"
+        );
+
+        let merge_sort =
+            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
+        let merge_projection = ProjectionExec::try_new(
+            vec![
+                (physical_col("ts", schema.as_ref())?, "ts".to_string()),
+                (physical_col("tag", schema.as_ref())?, "tag".to_string()),
+            ],
+            Arc::clone(&merge_sort),
+        )?;
+        let swapped_merge_sort = merge_sort
+            .try_swapping_with_projection(&merge_projection)?
+            .expect("MergeSortExec should accept the same projection swap as SPM");
+
+        assert!(
+            swapped_merge_sort
+                .as_any()
+                .downcast_ref::<MergeSortExec>()
+                .is_some(),
+            "MergeSortExec must rewrap projection swaps as MergeSortExec"
+        );
+        assert!(
+            swapped_merge_sort
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_none(),
+            "MergeSortExec must not expose a bare SPM after projection swap"
+        );
+        assert_eq!(swapped_merge_sort.fetch(), Some(1));
+        assert!(
+            swapped_merge_sort.children()[0]
+                .as_any()
+                .downcast_ref::<ProjectionExec>()
+                .is_some(),
+            "the projection should move below MergeSortExec"
+        );
+        let swapped_schema = swapped_merge_sort.schema();
+        assert_eq!(
+            swapped_schema
+                .fields()
+                .iter()
+                .map(|field| field.name().as_str())
+                .collect::<Vec<_>>(),
+            vec!["ts", "tag"],
+            "swapped MergeSortExec should expose the projected schema"
+        );
+
+        let projected_ordering = LexOrdering::new([PhysicalSortExpr::new(
+            physical_col("ts", swapped_merge_sort.children()[0].schema().as_ref())?,
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )])
+        .unwrap();
+        assert_eq!(
+            swapped_merge_sort.required_input_ordering(),
+            vec![Some(OrderingRequirements::from(projected_ordering))],
+            "projection swap must rewrite the ordering to the child projection's schema"
+        );
+        assert_eq!(
+            swapped_merge_sort.required_input_ordering(),
+            swapped_spm.required_input_ordering(),
+            "MergeSortExec projection swap should mirror SPM's ordering rewrite"
+        );
+
+        let spm_projection_without_sort_key = ProjectionExec::try_new(
+            vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())],
+            Arc::clone(&bare_spm),
+        )?;
+        let merge_projection_without_sort_key = ProjectionExec::try_new(
+            vec![(physical_col("tag", schema.as_ref())?, "tag".to_string())],
+            Arc::clone(&merge_sort),
+        )?;
+        assert!(
+            bare_spm
+                .try_swapping_with_projection(&spm_projection_without_sort_key)?
+                .is_none(),
+            "SPM must reject projection swaps that drop the sort key"
+        );
+        assert!(
+            merge_sort
+                .try_swapping_with_projection(&merge_projection_without_sort_key)?
+                .is_none(),
+            "MergeSortExec should reject the same projection swap as SPM"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn enforce_sorting_rewrite_keeps_merge_sort_exec_opaque() {
+        let schema = Arc::new(Schema::new(vec![Field::new("ts", DataType::Int64, false)]));
+        let input = Arc::new(EmptyExec::new(schema.clone()).with_partitions(2)) as _;
+        let ordering = test_ordering(schema.as_ref());
+
+        let bare_spm = Arc::new(
+            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(&input)).with_fetch(Some(1)),
+        ) as Arc<dyn ExecutionPlan>;
+        let optimized_spm = plan_with_order_breaking_variants(OrderPreservationContext::new(
+            bare_spm,
+            false,
+            vec![OrderPreservationContext::new(
+                Arc::clone(&input),
+                false,
+                vec![],
+            )],
+        ))
+        .unwrap()
+        .plan;
+        assert!(
+            optimized_spm
+                .as_any()
+                .downcast_ref::<CoalescePartitionsExec>()
+                .is_some(),
+            "this regression test must exercise EnforceSorting's bare SPM -> CoalescePartitionsExec rewrite"
+        );
+
+        let merge_sort =
+            Arc::new(MergeSortExec::new(ordering, input, Some(1))) as Arc<dyn ExecutionPlan>;
+        let optimized_merge_sort =
+            plan_with_order_breaking_variants(OrderPreservationContext::new(
+                Arc::clone(&merge_sort),
+                false,
+                vec![OrderPreservationContext::new(
+                    Arc::clone(merge_sort.children()[0]),
+                    false,
+                    vec![],
+                )],
+            ))
+            .unwrap()
+            .plan;
+        assert!(
+            optimized_merge_sort
+                .as_any()
+                .downcast_ref::<MergeSortExec>()
+                .is_some(),
+            "MergeSortExec must stay opaque to the bare SPM rewrite"
+        );
+        assert!(
+            optimized_merge_sort
+                .as_any()
+                .downcast_ref::<CoalescePartitionsExec>()
+                .is_none(),
+            "MergeSortExec(fetch) is the required distributed TopK merge stage, not an unordered coalesce"
+        );
+        assert_eq!(optimized_merge_sort.fetch(), Some(1));
+    }
+}
@@ -17,6 +17,7 @@
 use std::sync::Arc;

 use ahash::HashMap;
+use arrow_schema::SortOptions;
 use async_trait::async_trait;
 use catalog::CatalogManagerRef;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
@@ -29,6 +30,7 @@ use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion_common::{DataFusionError, TableReference};
 use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode};
+use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
 use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes};
 use session::context::QueryContext;
 use snafu::{OptionExt, ResultExt};
@@ -39,27 +41,54 @@ use table::table_name::TableName;

 use crate::dist_plan::PredicateExtractor;
 use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan};
-use crate::dist_plan::merge_sort::MergeSortLogicalPlan;
+use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan};
 use crate::dist_plan::region_pruner::ConstraintPruner;
 use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu};
 use crate::region_query::RegionQueryHandlerRef;

-/// Planner for convert merge sort logical plan to physical plan
+/// Planner for converting merge sort logical plan to physical plan.
 ///
-/// it is currently a fallback to sort, and doesn't change the execution plan:
-/// `MergeSort(MergeScan) -> Sort(MergeScan) - to physical plan -> ...`
-/// It should be applied after `DistExtensionPlanner`
-///
-/// (Later when actually impl this merge sort)
-///
-/// We should ensure the number of partition is not smaller than the number of region at present. Otherwise this would result in incorrect output.
+/// `MergeSortExec` always represents the distributed merge stage. It declares
+/// the required input ordering to DataFusion, so `EnforceSorting` inserts a
+/// `SortExec` below it when the input `MergeScanExec` cannot preserve per-region
+/// ordering, for example when one output partition may merge multiple region
+/// streams.
 pub struct MergeSortExtensionPlanner {}

+impl MergeSortExtensionPlanner {
+    fn ordering(
+        session_state: &SessionState,
+        merge_sort: &MergeSortLogicalPlan,
+    ) -> Result<LexOrdering> {
+        let ordering = merge_sort
+            .expr
+            .iter()
+            .map(|sort_expr| {
+                let physical_expr = session_state
+                    .create_physical_expr(sort_expr.expr.clone(), merge_sort.input.schema())?;
+                Ok(PhysicalSortExpr::new(
+                    physical_expr,
+                    SortOptions {
+                        descending: !sort_expr.asc,
+                        nulls_first: sort_expr.nulls_first,
+                    },
+                ))
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        LexOrdering::new(ordering).ok_or_else(|| {
+            DataFusionError::Internal(
+                "Expect MergeSort to have non-empty sort expressions".to_string(),
+            )
+        })
+    }
+}
+
 #[async_trait]
 impl ExtensionPlanner for MergeSortExtensionPlanner {
    async fn plan_extension(
        &self,
-        planner: &dyn PhysicalPlanner,
+        _planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        _logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
@@ -73,28 +102,24 @@ impl ExtensionPlanner for MergeSortExtensionPlanner {
                    .downcast_ref::<MergeScanLogicalPlan>()
                    .is_some()
            {
-                let merge_scan_exec = physical_inputs
-                    .first()
-                    .and_then(|p| p.as_any().downcast_ref::<MergeScanExec>())
-                    .ok_or(DataFusionError::Internal(format!(
+                let input = physical_inputs.first().cloned().ok_or_else(|| {
+                    DataFusionError::Internal(
+                        "Expect MergeSort to have one physical input".to_string(),
+                    )
+                })?;
+                if input.as_any().downcast_ref::<MergeScanExec>().is_none() {
+                    return Err(DataFusionError::Internal(format!(
                        "Expect MergeSort's input is a MergeScanExec, found {:?}",
                        physical_inputs
-                    )))?;
-
-                let partition_cnt = merge_scan_exec.partition_count();
-                let region_cnt = merge_scan_exec.region_count();
-                // if partition >= region, we know that every partition stream of merge scan is ordered
-                // and we only need to do a merge sort, otherwise fallback to quick sort
-                let can_merge_sort = partition_cnt >= region_cnt;
-                if can_merge_sort {
-                    // TODO(discord9): use `SortPreservingMergeExec here`
+                    )));
                }
-                // for now merge sort only exist in logical plan, and have the same effect as `Sort`
-                // doesn't change the execution plan, this will change in the future
-                let ret = planner
-                    .create_physical_plan(&merge_sort.clone().into_sort(), session_state)
-                    .await?;
-                Ok(Some(ret))
+
+                let ordering = Self::ordering(session_state, merge_sort)?;
+                Ok(Some(Arc::new(MergeSortExec::new(
+                    ordering,
+                    input,
+                    merge_sort.fetch,
+                ))))
            } else {
                Ok(None)
            }
@@ -25,6 +25,8 @@ use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use datafusion_common::Result as DfResult;
 use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning};

+use crate::dist_plan::MergeSortExec;
+
 #[derive(Debug)]
 pub struct EnsureGlobalLimitForFetch;

@@ -149,7 +151,8 @@ fn provided_global_fetch(plan: &Arc<dyn ExecutionPlan>) -> Option<usize> {
    let fetch = plan.fetch()?;
    (plan.as_any().is::<GlobalLimitExec>()
        || plan.as_any().is::<CoalescePartitionsExec>()
-        || plan.as_any().is::<SortPreservingMergeExec>())
+        || plan.as_any().is::<SortPreservingMergeExec>()
+        || plan.as_any().is::<MergeSortExec>())
    .then_some(fetch)
 }

@@ -327,6 +330,36 @@ mod tests {
        assert_eq!(child.fetch(), Some(5));
    }

+    #[test]
+    fn keeps_filter_under_parent_merge_sort_fetch() {
+        let (input, ordering) = ordered_input();
+        let filter = filter_fetch(input, 1);
+        let merge = merge_sort_fetch(ordering, filter, 1);
+
+        let optimized =
+            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
+        let child = optimized.children()[0];
+
+        assert!(optimized.as_any().is::<MergeSortExec>());
+        assert!(child.as_any().is::<FilterExec>());
+    }
+
+    #[test]
+    fn adds_tighter_global_fetch_under_looser_merge_sort_fetch() {
+        let (input, ordering) = ordered_input();
+        let filter = filter_fetch(input, 5);
+        let merge = merge_sort_fetch(ordering, filter, 10);
+
+        let optimized =
+            EnsureGlobalLimitForFetch::optimize_plan(merge, ParentContext::default()).unwrap();
+        let child = optimized.children()[0];
+
+        assert!(optimized.as_any().is::<MergeSortExec>());
+        assert!(child.as_any().is::<SortPreservingMergeExec>());
+        assert_eq!(child.fetch(), Some(5));
+        assert!(child.children()[0].as_any().is::<FilterExec>());
+    }
+
    #[test]
    fn preserves_parent_ordering_requirement() {
        let (input, ordering) = ordered_input();
@@ -552,6 +585,14 @@ mod tests {
        )
    }

+    fn merge_sort_fetch(
+        ordering: LexOrdering,
+        input: Arc<dyn ExecutionPlan>,
+        fetch: usize,
+    ) -> Arc<dyn ExecutionPlan> {
+        Arc::new(MergeSortExec::new(ordering, input, Some(fetch)))
+    }
+
    fn hash_repartition(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
        let partitioning = Partitioning::Hash(vec![col("a", input.schema().as_ref()).unwrap()], 3);
        Arc::new(RepartitionExec::try_new(input, partitioning).unwrap())
@@ -28,6 +28,7 @@ use datatypes::arrow::array::{
 use frontend::error::Error;
 use frontend::instance::Instance;
 use operator::error::Error as OperatorError;
+use query::datafusion::QUERY_PARALLELISM_HINT;
 use rstest::rstest;
 use rstest_reuse::apply;
 use servers::error as server_error;
@@ -89,6 +90,98 @@ async fn test_create_database_and_insert_query(instance: Arc<dyn MockInstance>)
    }
 }

+#[tokio::test(flavor = "multi_thread")]
+async fn test_distributed_scalar_latest_with_query_parallelism_below_regions() {
+    common_telemetry::init_default_ut_logging();
+
+    let distributed = crate::tests::create_distributed_instance(
+        "test_distributed_scalar_latest_with_query_parallelism_below_regions",
+    )
+    .await;
+    let frontend = distributed.frontend();
+
+    execute_sql(
+        &frontend,
+        r#"
+CREATE TABLE cpu (
+    rack STRING NULL,
+    os STRING NULL,
+    usage_user BIGINT NULL,
+    greptime_timestamp TIMESTAMP(9) NOT NULL,
+    TIME INDEX (greptime_timestamp)
+)
+PARTITION ON COLUMNS (rack) (
+    rack < '2',
+    rack >= '2' AND rack < '4',
+    rack >= '4' AND rack < '6',
+    rack >= '6' AND rack < '8',
+    rack >= '8'
+)
+ENGINE = mito
+WITH (append_mode = 'true', sst_format = 'flat')
+"#,
+    )
+    .await;
+
+    execute_sql(
+        &frontend,
+        r#"
+INSERT INTO cpu VALUES
+    ('1', 'linux', 10, '2023-06-12 01:04:49'),
+    ('1', 'linux', 15, '2023-06-12 01:04:50'),
+    ('3', 'windows', 25, '2023-06-12 01:05:00'),
+    ('5', 'mac', 30, '2023-06-12 01:03:00'),
+    ('7', 'linux', 45, '2023-06-12 02:00:00'),
+    ('2', 'linux', 20, '2023-06-12 01:04:51'),
+    ('2', 'windows', 22, '2023-06-12 01:06:00'),
+    ('4', 'mac', 12, '2023-06-12 00:59:00'),
+    ('6', 'linux', 35, '2023-06-12 01:04:55'),
+    ('8', 'windows', 50, '2023-06-12 02:10:00')
+"#,
+    )
+    .await;
+
+    let latest_sql = r#"
+SELECT rack, os, greptime_timestamp
+FROM cpu
+WHERE greptime_timestamp = (
+    SELECT greptime_timestamp
+    FROM cpu
+    ORDER BY greptime_timestamp DESC
+    LIMIT 1
+)
+"#;
+
+    let result = execute_sql_with_query_parallelism(&frontend, latest_sql, 1)
+        .await
+        .data
+        .pretty_print()
+        .await;
+    assert_eq!(
+        result,
+        r#"+------+---------+---------------------+
+| rack | os      | greptime_timestamp  |
+------+---------+---------------------+
+| 8    | windows | 2023-06-12T02:10:00 |
+------+---------+---------------------+"#
+    );
+
+    let explain =
+        execute_sql_with_query_parallelism(&frontend, &format!("EXPLAIN {latest_sql}"), 1)
+            .await
+            .data
+            .pretty_print()
+            .await;
+    assert!(
+        explain.contains("SortExec"),
+        "query_parallelism=1 with five regions should insert SortExec below MergeSortExec; explain:\n{explain}"
+    );
+    assert!(
+        explain.contains("MergeSortExec"),
+        "query_parallelism=1 with five regions should keep the distributed merge stage stable; explain:\n{explain}"
+    );
+}
+
 #[apply(both_instances_cases)]
 async fn test_show_create_table(instance: Arc<dyn MockInstance>) {
    let frontend = instance.frontend();
@@ -2876,6 +2969,16 @@ async fn execute_sql(instance: &Arc<Instance>, sql: &str) -> Output {
    execute_sql_with(instance, sql, QueryContext::arc()).await
 }

+async fn execute_sql_with_query_parallelism(
+    instance: &Arc<Instance>,
+    sql: &str,
+    parallelism: usize,
+) -> Output {
+    let mut query_ctx = QueryContext::with_db_name(None);
+    query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string());
+    execute_sql_with(instance, sql, Arc::new(query_ctx)).await
+}
+
 async fn try_execute_sql(instance: &Arc<Instance>, sql: &str) -> server_error::Result<Output> {
    try_execute_sql_with(instance, sql, QueryContext::arc()).await
 }
@@ -30,7 +30,7 @@ explain SELECT * FROM multi_partitions_test_table WHERE ts > cast(1000000000 as
 |_|_Filter: multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)_|
 |_|_TableScan: multi_partitions_test_table, partial_filters=[multi_partitions_test_table.ts > TimestampMillisecond(1000000000, None)]_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -37,7 +37,7 @@ tql explain (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr
 |               |                 Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)                                                                         |
 |               |                   TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)]                                  |
 |               | ]]                                                                                                                                                                                                                                                                      |
-| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST]                                                                                                                                              |
+| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST]                                                                                                                                                        |
 |               |   MergeScanExec: REDACTED
 |               |                                                                                                                                                                                                                                                                         |
 +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -54,7 +54,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED
@@ -244,7 +244,7 @@ tql explain (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time
 |               |                 Filter: aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None) AND aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)                                                                         |
 |               |                   TableScan: aggr_optimize_not, partial_filters=[aggr_optimize_not.greptime_timestamp >= TimestampMillisecond(1752591744001, None), aggr_optimize_not.greptime_timestamp <= TimestampMillisecond(1752592164000, None)]                                  |
 |               | ]]                                                                                                                                                                                                                                                                      |
-| physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST]                                                                                                                          |
+| physical_plan | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST]                                                                                                                                    |
 |               |   MergeScanExec: REDACTED
 |               |                                                                                                                                                                                                                                                                         |
 +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -261,7 +261,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED
@@ -485,9 +485,10 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize
 |               |                 PromSeriesDivideExec: tags=["a", "b", "c", "d"]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           |
 |               |                   SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             |
 |               |                     MergeScanExec: REDACTED
-|               |     SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
-|               |       CooperativeExec                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
-|               |         MergeScanExec: REDACTED
+|               |     RepartitionExec: partitioning=REDACTED
+|               |       MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+|               |         SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    |
+|               |           MergeScanExec: REDACTED
 |               |                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           |
 +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

@@ -514,8 +515,9 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize
 |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED
 |_|_|_SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] REDACTED
 |_|_|_MergeScanExec: REDACTED
+|_|_|_RepartitionExec: partitioning=REDACTED
+|_|_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED
 |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED
-|_|_|_CooperativeExec REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_CooperativeExec REDACTED
@@ -974,7 +976,7 @@ EXPLAIN SELECT pk_col_1, pk_col_2, sum(val_col_1) FROM step_aggr_extended GROUP
 |_|_Aggregate: groupBy=[[step_aggr_extended.pk_col_1, step_aggr_extended.pk_col_2]], aggr=[[sum(step_aggr_extended.val_col_1)]] |
 |_|_TableScan: step_aggr_extended_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [pk_col_1@0 ASC NULLS LAST, pk_col_2@1 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -287,7 +287,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x;
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -326,7 +326,7 @@ EXPLAIN SELECT x, COUNT(*) AS c FROM (SELECT a AS x FROM t) sq GROUP BY x ORDER
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -397,7 +397,7 @@ EXPLAIN SELECT sq.x FROM (SELECT a AS x FROM t) sq ORDER BY sq.x;
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -433,7 +433,7 @@ EXPLAIN SELECT y FROM (SELECT x AS y FROM (SELECT a AS x FROM t) sq1) sq2 ORDER
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [y@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [y@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -471,7 +471,7 @@ EXPLAIN SELECT x, x + 1 AS y FROM (SELECT a AS x FROM t) sq ORDER BY x;
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -625,7 +625,7 @@ EXPLAIN SELECT x FROM (SELECT a AS x FROM t) sq ORDER BY x LIMIT 2;
 |_|_Projection: t.a AS x_|
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST], fetch=2_|
+| physical_plan | MergeSortExec: [x@0 ASC NULLS LAST], fetch=2_|
 |_|_SortExec: TopK(fetch=2), expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
@@ -240,7 +240,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] |
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -267,7 +267,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_ProjectionExec: expr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED
@@ -606,7 +606,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] |
 |_|_TableScan: t1_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -632,7 +632,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_ProjectionExec: expr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED
@@ -240,7 +240,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] |
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -267,7 +267,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_ProjectionExec: expr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST]@0 as ordered_host, last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST]@1 as last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]@2 as last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]] REDACTED
@@ -606,7 +606,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] |
 |_|_TableScan: t1_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -632,7 +632,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_ProjectionExec: expr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST]@0 as ordered_host, last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]@1 as last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED
@@ -84,6 +84,82 @@ group by
 | windows | 2        |
 +---------+----------+

+-- Query the row at the latest timestamp through a scalar subquery.
+-- Keep the distributed merge sort visible, but redact the optional local SortExec
+-- that appears when query parallelism is lower than the region count.
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n
+EXPLAIN SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
+-+-+
+| plan_type_| plan_|
+-+-+
+| logical_plan_| Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_|
+|_|_Inner Join: cpu.greptime_timestamp = __scalar_sq_1.greptime_timestamp_|
+|_|_Projection: cpu.rack, cpu.os, cpu.greptime_timestamp_|
+|_|_MergeScan [is_placeholder=false, remote_input=[_|
+|_| TableScan: cpu_|
+|_| ]]_|
+|_|_SubqueryAlias: __scalar_sq_1_|
+|_|_Limit: skip=0, fetch=1_|
+|_|_MergeSort: cpu.greptime_timestamp DESC NULLS FIRST_|
+|_|_MergeScan [is_placeholder=false, remote_input=[_|
+|_| Limit: skip=0, fetch=1_|
+|_|_Sort: cpu.greptime_timestamp DESC NULLS FIRST_|
+|_|_Projection: cpu.greptime_timestamp_|
+|_|_TableScan: cpu_|
+|_| ]]_|
+| physical_plan | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(greptime_timestamp@0, greptime_timestamp@2)], projection=[rack@1, os@2, greptime_timestamp@3]_|
+|_|_MergeSortExec: [greptime_timestamp@0 DESC], fetch=1_|
+|_|_MergeScanExec: REDACTED
+|_|_ProjectionExec: expr=[rack@0 as rack, os@1 as os, greptime_timestamp@3 as greptime_timestamp]_|
+|_|_CooperativeExec_|
+|_|_MergeScanExec: REDACTED
+|_|_|
+-+-+
+
+SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
+------+---------+---------------------+
+| rack | os      | greptime_timestamp  |
+------+---------+---------------------+
+| 8    | windows | 2023-06-12T02:10:00 |
+------+---------+---------------------+
+
 drop table cpu;

 Affected Rows: 0
@@ -59,4 +59,47 @@ where
 group by
    os;

-drop table cpu;
+-- Query the row at the latest timestamp through a scalar subquery.
+-- Keep the distributed merge sort visible, but redact the optional local SortExec
+-- that appears when query parallelism is lower than the region count.
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE (?m)^\|_\|_SortExec:.*\n
+EXPLAIN SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
+SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
+drop table cpu;
@@ -296,7 +296,7 @@ explain analyze select tag from t where num > 6 order by ts desc limit 2;
 | stage | node | plan_|
 +-+-+-+
 | 0_| 0_|_ProjectionExec: expr=[tag@0 as tag] REDACTED
-|_|_|_SortPreservingMergeExec: [ts@1 DESC], fetch=2 REDACTED
+|_|_|_MergeSortExec: [ts@1 DESC], fetch=2 REDACTED
 |_|_|_SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[true] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
@@ -133,8 +133,8 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100);
 |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED
 |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED
 |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED
-|_|_|_SortExec: expr=[k@2 ASC, l@3 ASC, j@1 ASC], preserve_partitioning=[true] REDACTED
 |_|_|_RepartitionExec: partitioning=REDACTED
+|_|_|_MergeSortExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED
@@ -0,0 +1,187 @@
+CREATE TABLE netflow_raw (
+    greptime_timestamp TIMESTAMP TIME INDEX,
+    src STRING,
+    dst_port INT,
+    dst_addr STRING,
+    byte_count DOUBLE,
+    flow_id INT,
+    PRIMARY KEY (src, dst_port)
+);
+
+Affected Rows: 0
+
+INSERT INTO netflow_raw VALUES
+    (1000, 'source-a', 80, 'target-80-old', 100.0, 1),
+    (5000, 'source-a', 80, 'target-80-new', 500.0, 2),
+    (2000, 'source-a', 443, 'target-443-old', 200.0, 3),
+    (7000, 'source-a', 443, 'target-443-new', 700.0, 4),
+    (3000, 'source-a', 22, 'target-22-only', 300.0, 5),
+    (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6),
+    (8000, 'source-b', 9999, 'target-other-src', 999.0, 7);
+
+Affected Rows: 7
+
+-- Latest row per dst_port using a window subquery. The window result must be
+-- filtered outside the subquery, not directly in the WHERE clause.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        ROW_NUMBER() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS rn
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE rn = 1
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+---------------------+----------+------------------+------------+---------+
+| greptime_timestamp  | dst_port | dst_addr         | byte_count | flow_id |
+---------------------+----------+------------------+------------+---------+
+| 1970-01-01T00:00:07 | 443      | target-443-new   | 700.0      | 4       |
+| 1970-01-01T00:00:05 | 80       | target-80-new    | 500.0      | 2       |
+| 1970-01-01T00:00:04 | 8080     | target-8080-only | 400.0      | 6       |
+---------------------+----------+------------------+------------+---------+
+
+-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY.
+-- Wrap it to apply the final "latest groups first" ordering and LIMIT.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id
+FROM (
+    SELECT DISTINCT ON (dst_port)
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+    ORDER BY dst_port, greptime_timestamp DESC
+) latest_per_port
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+---------------------+----------+------------------+------------+---------+
+| greptime_timestamp  | dst_port | dst_addr         | byte_count | flow_id |
+---------------------+----------+------------------+------------+---------+
+| 1970-01-01T00:00:07 | 443      | target-443-new   | 700.0      | 4       |
+| 1970-01-01T00:00:05 | 80       | target-80-new    | 500.0      | 2       |
+| 1970-01-01T00:00:04 | 8080     | target-8080-only | 400.0      | 6       |
+---------------------+----------+------------------+------------+---------+
+
+-- RANK returns all rows tied at the latest timestamp for each dst_port.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        RANK() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS r
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE r = 1
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+---------------------+----------+------------------+------------+---------+---+
+| greptime_timestamp  | dst_port | dst_addr         | byte_count | flow_id | r |
+---------------------+----------+------------------+------------+---------+---+
+| 1970-01-01T00:00:07 | 443      | target-443-new   | 700.0      | 4       | 1 |
+| 1970-01-01T00:00:05 | 80       | target-80-new    | 500.0      | 2       | 1 |
+| 1970-01-01T00:00:04 | 8080     | target-8080-only | 400.0      | 6       | 1 |
+---------------------+----------+------------------+------------+---------+---+
+
+-- Ordered aggregate variant when the caller can enumerate the columns to fetch.
+SELECT
+    dst_port,
+    last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts,
+    last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr,
+    last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count,
+    last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id
+FROM netflow_raw
+WHERE src = 'source-a'
+  AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+GROUP BY dst_port
+ORDER BY latest_ts DESC, dst_port
+LIMIT 3;
+
+----------+---------------------+------------------+-------------------+----------------+
+| dst_port | latest_ts           | latest_dst_addr  | latest_byte_count | latest_flow_id |
+----------+---------------------+------------------+-------------------+----------------+
+| 443      | 1970-01-01T00:00:07 | target-443-new   | 700.0             | 4              |
+| 80       | 1970-01-01T00:00:05 | target-80-new    | 500.0             | 2              |
+| 8080     | 1970-01-01T00:00:04 | target-8080-only | 400.0             | 6              |
+----------+---------------------+------------------+-------------------+----------------+
+
+CREATE TABLE netflow_raw_ties (
+    greptime_timestamp TIMESTAMP TIME INDEX,
+    src STRING,
+    dst_port INT,
+    dst_addr STRING,
+    byte_count DOUBLE,
+    flow_id INT,
+    PRIMARY KEY (src, dst_port, flow_id)
+);
+
+Affected Rows: 0
+
+INSERT INTO netflow_raw_ties VALUES
+    (1000, 'source-a', 443, 'target-443-old', 100.0, 1),
+    (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2),
+    (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3),
+    (8000, 'source-a', 80, 'target-80-new', 800.0, 4),
+    (9500, 'source-b', 443, 'target-other-src', 950.0, 5);
+
+Affected Rows: 5
+
+-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        RANK() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS r
+    FROM netflow_raw_ties
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE r = 1
+ORDER BY greptime_timestamp DESC, dst_port, flow_id;
+
+---------------------+----------+------------------+------------+---------+---+
+| greptime_timestamp  | dst_port | dst_addr         | byte_count | flow_id | r |
+---------------------+----------+------------------+------------+---------+---+
+| 1970-01-01T00:00:09 | 443      | target-443-tie-a | 900.0      | 2       | 1 |
+| 1970-01-01T00:00:09 | 443      | target-443-tie-b | 901.0      | 3       | 1 |
+| 1970-01-01T00:00:08 | 80       | target-80-new    | 800.0      | 4       | 1 |
+---------------------+----------+------------------+------------+---------+---+
+
+DROP TABLE netflow_raw_ties;
+
+Affected Rows: 0
+
+DROP TABLE netflow_raw;
+
+Affected Rows: 0
+
@@ -0,0 +1,134 @@
+CREATE TABLE netflow_raw (
+    greptime_timestamp TIMESTAMP TIME INDEX,
+    src STRING,
+    dst_port INT,
+    dst_addr STRING,
+    byte_count DOUBLE,
+    flow_id INT,
+    PRIMARY KEY (src, dst_port)
+);
+
+INSERT INTO netflow_raw VALUES
+    (1000, 'source-a', 80, 'target-80-old', 100.0, 1),
+    (5000, 'source-a', 80, 'target-80-new', 500.0, 2),
+    (2000, 'source-a', 443, 'target-443-old', 200.0, 3),
+    (7000, 'source-a', 443, 'target-443-new', 700.0, 4),
+    (3000, 'source-a', 22, 'target-22-only', 300.0, 5),
+    (4000, 'source-a', 8080, 'target-8080-only', 400.0, 6),
+    (8000, 'source-b', 9999, 'target-other-src', 999.0, 7);
+
+-- Latest row per dst_port using a window subquery. The window result must be
+-- filtered outside the subquery, not directly in the WHERE clause.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        ROW_NUMBER() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS rn
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE rn = 1
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+-- DISTINCT ON keeps the first row per dst_port according to its own ORDER BY.
+-- Wrap it to apply the final "latest groups first" ordering and LIMIT.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id
+FROM (
+    SELECT DISTINCT ON (dst_port)
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+    ORDER BY dst_port, greptime_timestamp DESC
+) latest_per_port
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+-- RANK returns all rows tied at the latest timestamp for each dst_port.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        RANK() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS r
+    FROM netflow_raw
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE r = 1
+ORDER BY greptime_timestamp DESC, dst_port
+LIMIT 3;
+
+-- Ordered aggregate variant when the caller can enumerate the columns to fetch.
+SELECT
+    dst_port,
+    last_value(greptime_timestamp ORDER BY greptime_timestamp) AS latest_ts,
+    last_value(dst_addr ORDER BY greptime_timestamp) AS latest_dst_addr,
+    last_value(byte_count ORDER BY greptime_timestamp) AS latest_byte_count,
+    last_value(flow_id ORDER BY greptime_timestamp) AS latest_flow_id
+FROM netflow_raw
+WHERE src = 'source-a'
+  AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+GROUP BY dst_port
+ORDER BY latest_ts DESC, dst_port
+LIMIT 3;
+
+CREATE TABLE netflow_raw_ties (
+    greptime_timestamp TIMESTAMP TIME INDEX,
+    src STRING,
+    dst_port INT,
+    dst_addr STRING,
+    byte_count DOUBLE,
+    flow_id INT,
+    PRIMARY KEY (src, dst_port, flow_id)
+);
+
+INSERT INTO netflow_raw_ties VALUES
+    (1000, 'source-a', 443, 'target-443-old', 100.0, 1),
+    (9000, 'source-a', 443, 'target-443-tie-a', 900.0, 2),
+    (9000, 'source-a', 443, 'target-443-tie-b', 901.0, 3),
+    (8000, 'source-a', 80, 'target-80-new', 800.0, 4),
+    (9500, 'source-b', 443, 'target-other-src', 950.0, 5);
+
+-- RANK with r = 1 keeps all rows tied at the latest timestamp per dst_port.
+SELECT greptime_timestamp, dst_port, dst_addr, byte_count, flow_id, r
+FROM (
+    SELECT
+        greptime_timestamp,
+        dst_port,
+        dst_addr,
+        byte_count,
+        flow_id,
+        RANK() OVER (
+            PARTITION BY dst_port
+            ORDER BY greptime_timestamp DESC
+        ) AS r
+    FROM netflow_raw_ties
+    WHERE src = 'source-a'
+      AND greptime_timestamp >= '1970-01-01T00:00:02'::timestamp
+) latest_per_port
+WHERE r = 1
+ORDER BY greptime_timestamp DESC, dst_port, flow_id;
+
+DROP TABLE netflow_raw_ties;
+
+DROP TABLE netflow_raw;
@@ -240,7 +240,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t.host]], aggr=[[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] |
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -267,7 +267,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
@@ -603,7 +603,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t1.host]], aggr=[[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] |
 |_|_TableScan: t1_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -629,7 +629,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
@@ -240,7 +240,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t.host]], aggr=[[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.not_pk) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST]]] |
 |_|_TableScan: t_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -267,7 +267,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
@@ -603,7 +603,7 @@ order by ordered_host;
 |_|_Aggregate: groupBy=[[t1.host]], aggr=[[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST]]] |
 |_|_TableScan: t1_|
 |_| ]]_|
-| physical_plan | SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST]_|
+| physical_plan | MergeSortExec: [ordered_host@0 ASC NULLS LAST]_|
 |_|_MergeScanExec: REDACTED
 |_|_|
 +-+-+
@@ -629,7 +629,7 @@ order by ordered_host;
 +-+-+-+
 | stage | node | plan_|
 +-+-+-+
-| 0_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
+| 0_| 0_|_MergeSortExec: [ordered_host@0 ASC NULLS LAST] REDACTED
 |_|_|_MergeScanExec: REDACTED
 |_|_|_|
 | 1_| 0_|_SortPreservingMergeExec: [ordered_host@0 ASC NULLS LAST] REDACTED
@@ -197,6 +197,31 @@ group by
 | windows | 2        |
 +---------+----------+

+-- Query the row at the latest timestamp through a scalar subquery.
+SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
+------+---------+---------------------+
+| rack | os      | greptime_timestamp  |
+------+---------+---------------------+
+| 8    | windows | 2023-06-12T02:10:00 |
+------+---------+---------------------+
+
 drop table cpu;

 Affected Rows: 0
@@ -120,4 +120,23 @@ where
 group by
    os;

+-- Query the row at the latest timestamp through a scalar subquery.
+SELECT
+    rack,
+    os,
+    greptime_timestamp
+FROM
+    cpu
+WHERE
+    greptime_timestamp = (
+        SELECT
+            greptime_timestamp
+        FROM
+            cpu
+        ORDER BY
+            greptime_timestamp DESC
+        LIMIT
+            1
+    );
+
 drop table cpu;