fix: preserve local TopK limits during exact sort pushdown

datafusion/physical-optimizer/src/pushdown_sort.rs

@@ -167,11 +167,15 @@ impl PhysicalOptimizerRule for PushdownSort {
                     //
                     // If the SortExec carried a fetch (LIMIT), we must preserve it.
                     // First try pushing the limit into the source via `with_fetch()`.
-                    // If the source doesn't support `with_fetch`, fall back to
-                    // wrapping with GlobalLimitExec.
+                    // If the source doesn't support `with_fetch`, preserve whether
+                    // the limit applies to each partition or to the entire input.
                     if let Some(fetch) = sort_exec.fetch() {
                         let inner = inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
-                            Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch)))
+                            if sort_exec.preserve_partitioning() {
+                                Arc::new(LocalLimitExec::new(inner, fetch))
+                            } else {
+                                Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch)))
+                            }
                         });
                         Ok(Transformed::yes(inner))
                     } else {

datafusion/sqllogictest/src/test_context.rs

@@ -73,6 +73,7 @@ use log::info;
 use sqlparser::ast;
 use tempfile::TempDir;
 
+mod exact_sort_pushdown;
 mod range_partitioning;
 
 /// Context for running tests
@@ -182,6 +183,10 @@ impl TestContext {
 
         let file_name = relative_path.file_name().unwrap().to_str().unwrap();
         match file_name {
+            "sort_pushdown.slt" => {
+                exact_sort_pushdown::register_exact_sort_source(test_ctx.session_ctx())
+                    .expect("register exact sort source");
+            }
             "parquet_missing_bounds.slt" => {
                 register_parquet_missing_bounds(&mut test_ctx).await;
             }

datafusion/sqllogictest/src/test_context/exact_sort_pushdown.rs

@@ -0,0 +1,141 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::fmt::{self, Formatter};
+use std::sync::Arc;
+
+use arrow::array::record_batch;
+use arrow::datatypes::{self as arrow_schema, SchemaRef};
+use async_trait::async_trait;
+use datafusion::catalog::{Session, TableProvider};
+use datafusion::common::Result;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::logical_expr::{Expr, TableType};
+use datafusion::physical_expr::expressions::Column;
+use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
+use datafusion::physical_plan::test::TestMemoryExec;
+use datafusion::physical_plan::{
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult,
+};
+use datafusion::prelude::SessionContext;
+
+pub(super) fn register_exact_sort_source(ctx: &SessionContext) -> Result<()> {
+    let partitions = (0..4)
+        .map(|i| {
+            Ok(vec![record_batch!(
+                ("group_id", Int32, [0, 0, 1, 1]),
+                ("value", Int32, [i, i + 4, i + 8, i + 12])
+            )?])
+        })
+        .collect::<Result<Vec<_>>>()?;
+    let input = TestMemoryExec::try_new(&partitions, partitions[0][0].schema(), None)?;
+    ctx.register_table("exact_sort_source", Arc::new(ExactSortSource(input)))?;
+    Ok(())
+}
+
+fn ordering() -> LexOrdering {
+    [
+        PhysicalSortExpr::new_default(Arc::new(Column::new("group_id", 0))).nulls_last(),
+        PhysicalSortExpr::new_default(Arc::new(Column::new("value", 1))).nulls_last(),
+    ]
+    .into()
+}
+
+/// Accepts exact sort pushdown without supporting fetch pushdown. Its rows are
+/// already sorted within each partition; the hook makes that ordering explicit.
+#[derive(Debug, Clone)]
+struct ExactSortSource(TestMemoryExec);
+
+#[async_trait]
+impl TableProvider for ExactSortSource {
+    fn schema(&self) -> SchemaRef {
+        self.0.schema()
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        _: &dyn Session,
+        projection: Option<&[usize]>,
+        _: &[Expr],
+        _: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        assert!(projection.is_none_or(|indices| indices == [0, 1]));
+        Ok(Arc::new(self.clone()))
+    }
+}
+
+impl DisplayAs for ExactSortSource {
+    fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
+        write!(f, "ExactSortSource")
+    }
+}
+
+impl ExecutionPlan for ExactSortSource {
+    fn name(&self) -> &str {
+        "ExactSortSource"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        self.0.properties()
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        assert!(children.is_empty());
+        Ok(self)
+    }
+
+    fn apply_expressions(
+        &self,
+        _: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        self.0.execute(partition, context)
+    }
+
+    fn try_pushdown_sort(
+        &self,
+        order: &[PhysicalSortExpr],
+    ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
+        if order != ordering().as_ref() {
+            return Ok(SortOrderPushdownResult::Unsupported);
+        }
+        Ok(SortOrderPushdownResult::Exact {
+            inner: Arc::new(Self(
+                self.0.clone().try_with_sort_information(vec![ordering()])?,
+            )),
+        })
+    }
+}

datafusion/sqllogictest/test_files/sort_pushdown.slt

@@ -22,6 +22,54 @@ SET datafusion.execution.parquet.pushdown_filters = true;
 statement ok
 SET datafusion.optimizer.enable_sort_pushdown = true;
 
+# exact_sort_source has four partitions with interleaved values 0..15.
+# Values 0..7 belong to group 0; values 8..15 belong to group 1.
+# It supports exact (group_id, value) ordering, but cannot apply a limit itself.
+#
+# The partitioned window requires redistributing this source's rows by group.
+# The outer LIMIT lets the optimizer insert a local TopK before redistribution,
+# exercising the standalone sort-pushdown path. Without the window, this query
+# uses the already-correct merge/sort path.
+# Removing the source-side sort must retain its per-partition limit.
+query III
+SELECT group_id, value,
+       SUM(value) OVER (
+           PARTITION BY group_id ORDER BY value ROWS UNBOUNDED PRECEDING
+       ) AS running_sum
+FROM exact_sort_source
+ORDER BY group_id, value
+LIMIT 3;
+----
+0 0 0
+0 1 1
+0 2 3
+
+# The pushed-down sort must leave a LocalLimitExec below the repartition.
+statement ok
+SET datafusion.explain.physical_plan_only = true;
+
+query TT
+EXPLAIN
+SELECT group_id, value,
+       SUM(value) OVER (
+           PARTITION BY group_id ORDER BY value ROWS UNBOUNDED PRECEDING
+       ) AS running_sum
+FROM exact_sort_source
+ORDER BY group_id, value
+LIMIT 3;
+----
+physical_plan
+01)SortPreservingMergeExec: [group_id@0 ASC NULLS LAST, value@1 ASC NULLS LAST], fetch=3
+02)--ProjectionExec: expr=[group_id@0 as group_id, value@1 as value, sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as running_sum]
+03)----BoundedWindowAggExec: wdw=[sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+04)------SortExec: TopK(fetch=3), expr=[group_id@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true]
+05)--------RepartitionExec: partitioning=Hash([group_id@0], 4), input_partitions=4
+06)----------LocalLimitExec: fetch=3
+07)------------ExactSortSource
+
+statement ok
+RESET datafusion.explain.physical_plan_only;
+
 # Test 1: Sort Pushdown for ordered Parquet files
 # Create a sorted dataset
 statement ok