You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/25 16:10:14 UTC
(datafusion) branch main updated: Fix PushdownSort dropping LIMIT when eliminating SortExec (#21744)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6f1040bac3 Fix PushdownSort dropping LIMIT when eliminating SortExec (#21744)
6f1040bac3 is described below
commit 6f1040bac3e549c2fd4bb69e17c50987ca6c5fbf
Author: Sergei Grebnov <se...@gmail.com>
AuthorDate: Sat Apr 25 19:08:07 2026 +0300
Fix PushdownSort dropping LIMIT when eliminating SortExec (#21744)
## Which issue does this PR close?
When `PushdownSort` removes a `SortExec` because a source returns
`Exact` (guaranteeing ordering), any `fetch` (LIMIT) on the `SortExec`
is silently dropped if the underlying plan does not support
`with_fetch()`.
For example, `ProjectionExec` supports `try_pushdown_sort` (delegating
to its child) but does not implement `with_fetch()`. A plan like
`SortExec(fetch=10) → ProjectionExec → source` that gets sort-eliminated
loses the limit.
## What changes are included in this PR?
In the `Exact` branch of `PushdownSort`, when the eliminated `SortExec`
carried a `fetch`:
1. Try `with_fetch()` on the pushed-down source first
2. If `with_fetch()` returns `None`, fall back to wrapping with
`GlobalLimitExec`
## Are these changes tested?
Yes. Three new unit tests:
- `test_sort_pushdown_exact_no_fetch_no_limit` — Exact elimination
without fetch: no limit wrapper added
- `test_sort_pushdown_exact_preserves_fetch_with_global_limit` — Exact
elimination with fetch, source does NOT support `with_fetch()`:
`GlobalLimitExec` wrapper added
- `test_sort_pushdown_exact_preserves_fetch_with_source_support` — Exact
elimination with fetch, source supports `with_fetch()`: limit pushed
into source directly
## Are there any user-facing changes?
No.
---
.../core/tests/physical_optimizer/pushdown_sort.rs | 96 +++++++++++++++++++++-
.../core/tests/physical_optimizer/test_utils.rs | 85 ++++++++++++++++---
datafusion/physical-optimizer/src/pushdown_sort.rs | 26 ++++--
3 files changed, 185 insertions(+), 22 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs
index 2f66aa3c3e..e2700c3174 100644
--- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs
+++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs
@@ -32,10 +32,10 @@ use datafusion_physical_optimizer::pushdown_sort::PushdownSort;
use std::sync::Arc;
use crate::physical_optimizer::test_utils::{
- OptimizationTest, coalesce_partitions_exec, parquet_exec, parquet_exec_with_sort,
- projection_exec, projection_exec_with_alias, repartition_exec, schema,
- simple_projection_exec, sort_exec, sort_exec_with_fetch, sort_expr, sort_expr_named,
- test_scan_with_ordering,
+ OptimizationTest, TestScan, coalesce_partitions_exec, parquet_exec,
+ parquet_exec_with_sort, projection_exec, projection_exec_with_alias,
+ repartition_exec, schema, simple_projection_exec, sort_exec, sort_exec_with_fetch,
+ sort_expr, sort_expr_named, test_scan_with_ordering,
};
#[test]
@@ -996,3 +996,91 @@ fn test_sort_pushdown_with_test_scan_arbitrary_ordering() {
"
);
}
+
+// ============================================================================
+// EXACT PUSHDOWN TESTS (source guarantees ordering, SortExec removed)
+// ============================================================================
+
+#[test]
+fn test_sort_pushdown_exact_no_fetch_no_limit() {
+ // When a source returns Exact (without fetch), the SortExec should be
+ // removed entirely with no GlobalLimitExec wrapper.
+ let schema = schema();
+ let a = sort_expr("a", &schema);
+ let b = sort_expr("b", &schema);
+ let source =
+ Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true));
+
+ let ordering = LexOrdering::new(vec![a, b.reverse()]).unwrap();
+ let plan = sort_exec(ordering, source);
+
+ insta::assert_snapshot!(
+ OptimizationTest::new(plan, PushdownSort::new(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: expr=[a@0 ASC, b@1 DESC NULLS LAST], preserve_partitioning=[false]
+ - TestScan
+ output:
+ Ok:
+ - TestScan: requested_ordering=[a@0 ASC, b@1 DESC NULLS LAST]
+ "
+ );
+}
+
+#[test]
+fn test_sort_pushdown_exact_preserves_fetch_with_global_limit() {
+ // When a source returns Exact but does NOT support with_fetch(),
+ // the optimizer must wrap the result with GlobalLimitExec to preserve
+ // the LIMIT from the eliminated SortExec.
+ let schema = schema();
+ let a = sort_expr("a", &schema);
+ let source =
+ Arc::new(TestScan::new(schema.clone(), vec![]).with_exact_pushdown(true));
+
+ let ordering = LexOrdering::new(vec![a]).unwrap();
+ let plan = sort_exec_with_fetch(ordering, Some(10), source);
+
+ insta::assert_snapshot!(
+ OptimizationTest::new(plan, PushdownSort::new(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[false]
+ - TestScan
+ output:
+ Ok:
+ - GlobalLimitExec: skip=0, fetch=10
+ - TestScan: requested_ordering=[a@0 ASC]
+ "
+ );
+}
+
+#[test]
+fn test_sort_pushdown_exact_preserves_fetch_with_source_support() {
+ // When a source returns Exact AND supports with_fetch(),
+ // the limit should be pushed into the source directly (no GlobalLimitExec).
+ let schema = schema();
+ let a = sort_expr("a", &schema);
+ let source = Arc::new(
+ TestScan::new(schema.clone(), vec![])
+ .with_exact_pushdown(true)
+ .with_supports_fetch(true),
+ );
+
+ let ordering = LexOrdering::new(vec![a]).unwrap();
+ let plan = sort_exec_with_fetch(ordering, Some(10), source);
+
+ insta::assert_snapshot!(
+ OptimizationTest::new(plan, PushdownSort::new(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[false]
+ - TestScan
+ output:
+ Ok:
+ - TestScan: requested_ordering=[a@0 ASC], fetch=10
+ "
+ );
+}
diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs
index 4b6db1abc4..6814ab2358 100644
--- a/datafusion/core/tests/physical_optimizer/test_utils.rs
+++ b/datafusion/core/tests/physical_optimizer/test_utils.rs
@@ -830,8 +830,17 @@ pub fn sort_expr_named(name: &str, index: usize) -> PhysicalSortExpr {
}
}
-/// A test data source that can display any requested ordering
-/// This is useful for testing sort pushdown behavior
+/// A test data source that can display any requested ordering.
+/// This is useful for testing sort pushdown behavior.
+///
+/// ## Configuration
+///
+/// - `exact_pushdown`: if `true`, `try_pushdown_sort` returns `Exact`
+/// (source guarantees ordering, SortExec can be removed); if `false`
+/// (default), returns `Inexact` (SortExec kept).
+/// - `supports_fetch`: if `true`, `with_fetch()` returns `Some` so the
+/// optimizer can push a LIMIT into the source; if `false` (default),
+/// `with_fetch()` returns `None`, forcing a `GlobalLimitExec` wrapper.
#[derive(Debug, Clone)]
pub struct TestScan {
schema: SchemaRef,
@@ -839,6 +848,12 @@ pub struct TestScan {
plan_properties: Arc<PlanProperties>,
// Store the requested ordering for display
requested_ordering: Option<LexOrdering>,
+ /// If true, `try_pushdown_sort` returns `Exact` instead of `Inexact`.
+ exact_pushdown: bool,
+ /// If true, `with_fetch()` returns `Some(...)` (source absorbs the limit).
+ supports_fetch: bool,
+ /// The fetch (LIMIT) value pushed into this scan via `with_fetch()`.
+ fetch: Option<usize>,
}
impl TestScan {
@@ -872,6 +887,9 @@ impl TestScan {
output_ordering,
plan_properties: Arc::new(plan_properties),
requested_ordering: None,
+ exact_pushdown: false,
+ supports_fetch: false,
+ fetch: None,
}
}
@@ -879,6 +897,18 @@ impl TestScan {
pub fn with_ordering(schema: SchemaRef, ordering: LexOrdering) -> Self {
Self::new(schema, vec![ordering])
}
+
+ /// Set whether `try_pushdown_sort` returns `Exact` (true) or `Inexact` (false).
+ pub fn with_exact_pushdown(mut self, exact: bool) -> Self {
+ self.exact_pushdown = exact;
+ self
+ }
+
+ /// Set whether `with_fetch()` returns `Some` (true) or `None` (false).
+ pub fn with_supports_fetch(mut self, supports: bool) -> Self {
+ self.supports_fetch = supports;
+ self
+ }
}
impl DisplayAs for TestScan {
@@ -886,9 +916,9 @@ impl DisplayAs for TestScan {
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
write!(f, "TestScan")?;
+ let mut sep = ": ";
if !self.output_ordering.is_empty() {
- write!(f, ": output_ordering=[")?;
- // Format the ordering in a readable way
+ write!(f, "{sep}output_ordering=[")?;
for (i, sort_expr) in self.output_ordering[0].iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
@@ -896,10 +926,10 @@ impl DisplayAs for TestScan {
write!(f, "{sort_expr}")?;
}
write!(f, "]")?;
+ sep = ", ";
}
- // This is the key part - show what ordering was requested
if let Some(ref req) = self.requested_ordering {
- write!(f, ", requested_ordering=[")?;
+ write!(f, "{sep}requested_ordering=[")?;
for (i, sort_expr) in req.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
@@ -907,6 +937,10 @@ impl DisplayAs for TestScan {
write!(f, "{sort_expr}")?;
}
write!(f, "]")?;
+ sep = ", ";
+ }
+ if let Some(fetch) = self.fetch {
+ write!(f, "{sep}fetch={fetch}")?;
}
Ok(())
}
@@ -953,6 +987,20 @@ impl ExecutionPlan for TestScan {
Ok(Arc::new(Statistics::new_unknown(&self.schema)))
}
+ fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+ if self.supports_fetch {
+ let mut new_scan = self.clone();
+ new_scan.fetch = fetch;
+ Some(Arc::new(new_scan))
+ } else {
+ None
+ }
+ }
+
+ fn fetch(&self) -> Option<usize> {
+ self.fetch
+ }
+
// This is the key method - implement sort pushdown
fn try_pushdown_sort(
&self,
@@ -965,10 +1013,27 @@ impl ExecutionPlan for TestScan {
let mut new_scan = self.clone();
new_scan.requested_ordering = requested_ordering;
- // Always return Inexact to keep the Sort node (like Phase 1 behavior)
- Ok(SortOrderPushdownResult::Inexact {
- inner: Arc::new(new_scan),
- })
+ if self.exact_pushdown {
+ // Update plan properties to reflect the guaranteed ordering
+ let orderings: Vec<Vec<PhysicalSortExpr>> = vec![order.to_vec()];
+ let eq_properties = EquivalenceProperties::new_with_orderings(
+ Arc::clone(&self.schema),
+ orderings,
+ );
+ new_scan.plan_properties = Arc::new(PlanProperties::new(
+ eq_properties,
+ Partitioning::UnknownPartitioning(1),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+ Ok(SortOrderPushdownResult::Exact {
+ inner: Arc::new(new_scan),
+ })
+ } else {
+ Ok(SortOrderPushdownResult::Inexact {
+ inner: Arc::new(new_scan),
+ })
+ }
}
fn apply_expressions(
diff --git a/datafusion/physical-optimizer/src/pushdown_sort.rs b/datafusion/physical-optimizer/src/pushdown_sort.rs
index 9da8ce60e1..40a6fe2c20 100644
--- a/datafusion/physical-optimizer/src/pushdown_sort.rs
+++ b/datafusion/physical-optimizer/src/pushdown_sort.rs
@@ -61,6 +61,7 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::SortOrderPushdownResult;
use datafusion_physical_plan::buffer::BufferExec;
+use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use std::sync::Arc;
@@ -104,8 +105,12 @@ impl PhysicalOptimizerRule for PushdownSort {
let required_ordering = sort_child.expr();
match sort_input.try_pushdown_sort(required_ordering)? {
SortOrderPushdownResult::Exact { inner } => {
+ // Preserve fetch (LIMIT) from the eliminated SortExec.
+ // Use LocalLimitExec (not Global) since input is multi-partition.
let inner = if let Some(fetch) = sort_child.fetch() {
- inner.with_fetch(Some(fetch)).unwrap_or(inner)
+ inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
+ Arc::new(LocalLimitExec::new(inner, fetch))
+ })
} else {
inner
};
@@ -149,14 +154,19 @@ impl PhysicalOptimizerRule for PushdownSort {
match sort_input.try_pushdown_sort(required_ordering)? {
SortOrderPushdownResult::Exact { inner } => {
// Data source guarantees perfect ordering - remove the Sort operator.
- // Preserve the fetch (LIMIT) from the original SortExec so the
- // data source can stop reading early.
- let inner = if let Some(fetch) = sort_exec.fetch() {
- inner.with_fetch(Some(fetch)).unwrap_or(inner)
+ //
+ // 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 let Some(fetch) = sort_exec.fetch() {
+ let inner = inner.with_fetch(Some(fetch)).unwrap_or_else(|| {
+ Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch)))
+ });
+ Ok(Transformed::yes(inner))
} else {
- inner
- };
- Ok(Transformed::yes(inner))
+ Ok(Transformed::yes(inner))
+ }
}
SortOrderPushdownResult::Inexact { inner } => {
// Data source is optimized for the ordering but not perfectly sorted
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org