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/09/06 01:52:28 UTC
(datafusion) branch main updated: fix: prevent DISTINCT ORDER BY LIMIT from returning too few rows (#24932)
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 408527ae58 fix: prevent DISTINCT ORDER BY LIMIT from returning too few rows (#24932)
408527ae58 is described below
commit 408527ae58b5b706f5c02df8c37e096122cd8c78
Author: Geoffrey Claude <ge...@datadoghq.com>
AuthorDate: Sun Sep 6 01:41:53 2026 +0000
fix: prevent DISTINCT ORDER BY LIMIT from returning too few rows (#24932)
## Which issue does this PR close?
- Closes #24927.
## Rationale for this change
With multiple partitions, `SELECT DISTINCT ... ORDER BY ... LIMIT` can
return fewer rows than requested even when enough distinct values exist.
`DISTINCT` runs in partial and final stages. Sort pushdown was moving
the final TopK `fetch` (the sort's row limit) below the final
aggregation, while duplicate partial results still existed:
```text
SortPreservingMergeExec: fetch=3
AggregateExec: final DISTINCT
SortExec: TopK(fetch=3) <-- limit is applied too early
AggregateExec: partial DISTINCT
```
Duplicates could fill the limit and discard another value before the
final `DISTINCT` had a chance to remove them.
The corrected plan keeps the limit above the final `DISTINCT`:
```text
SortPreservingMergeExec: fetch=3 <-- limit is applied after DISTINCT
AggregateExec: final DISTINCT
SortExec <-- ordering only; no limit
AggregateExec: partial DISTINCT
```
I bisected this regression to
[`450c861a8c`](https://github.com/apache/datafusion/commit/450c861a8c38f8934b134368949606aaf3287792)
(#14821), first released in 47.0.0. That refactor made the
ordering-satisfied fast path forward `fetch` to the current plan's
children. It correctly established that the current plan preserved the
requested ordering, but preserving order does not mean that moving a
limit is safe. The shortcut consequently bypassed the existing
limit-pushdown and cardinality checks: the parent of that commit returns
`0`, `1`, and `2`, while the commit itself returns only `0` and `2`.
## What changes are included in this PR?
When a plan already satisfies the requested ordering, sort pushdown
first tries to keep a pending `fetch` on that plan through the generic
`ExecutionPlan::with_fetch` capability. It forwards the fetch to
children only when the plan explicitly supports limit pushdown and
produces exactly one output row for every input row.
If an already ordered plan can neither retain nor safely forward the
fetch, the optimizer keeps the limit above it with a `LocalLimitExec` or
`GlobalLimitExec`, adding a sort-preserving merge when a single output
partition is required. This carries the limit without introducing a
redundant sort and gives subsequent optimizer runs the same canonical
plan shape.
## What is the testing strategy for this PR?
An SQLLogicTest in `limit.slt` uses the exact two-partition query from
#24927 and verifies that it returns `0`, `1`, and `2`. Focused physical
optimizer tests verify the stable limit-carrier plan across another
optimizer pass and verify that fetch is still forwarded through a safe
equal-cardinality plan.
## Are there any user-facing changes?
Yes. Affected queries now return all requested distinct rows. There are
no API or configuration changes.
---
.../physical_optimizer/ensure_requirements.rs | 101 +++++++++++++++-
.../enforce_sorting/sort_pushdown.rs | 128 ++++++++++++++++++---
datafusion/sqllogictest/test_files/limit.slt | 27 +++++
3 files changed, 237 insertions(+), 19 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs
index 86b60519da..06b57a8630 100644
--- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs
+++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs
@@ -30,6 +30,9 @@ use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
use datafusion_physical_optimizer::ensure_requirements::enforce_sorting::{
PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
};
+use datafusion_physical_optimizer::ensure_requirements::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+};
use std::sync::Arc;
@@ -41,7 +44,7 @@ use datafusion_physical_expr::{
EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
};
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
-use datafusion_physical_plan::limit::GlobalLimitExec;
+use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::union::UnionExec;
use datafusion_physical_plan::{
@@ -835,6 +838,102 @@ fn test_idempotent_sort_aggregate_sort_aggregate() {
assert_idempotent(limit);
}
+/// Verifies the safe fallback for an already ordered, multi-partition plan that
+/// can neither retain a fetch nor pass it to its children. The fetch must stay
+/// above that plan in a `LocalLimitExec`, retain the ordering needed by the
+/// merge, and produce exactly the same plan on the next optimizer pass.
+///
+/// This prevents a late fallback from creating a temporary plan shape that
+/// `EnsureRequirements` rewrites when applied again.
+#[test]
+fn test_fetch_fallback_is_stable() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+ let union = UnionExec::try_new(vec![Arc::clone(&source), source]).unwrap();
+ let ordering = sort_expr_on("a", 0, false, false);
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
+ SortExec::new(ordering.clone(), union)
+ .with_fetch(Some(10))
+ .with_preserve_partitioning(true),
+ );
+ let merge: Arc<dyn ExecutionPlan> = Arc::new(
+ SortPreservingMergeExec::new(ordering.clone(), sort).with_fetch(Some(10)),
+ );
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ merge,
+ Some(OrderingRequirements::from(ordering.clone())),
+ Distribution::SinglePartition,
+ Some(10),
+ ));
+
+ // Exercise sort pushdown directly so the earlier sort-cleanup phase does
+ // not canonicalize the redundant TopK before it reaches this fallback.
+ let mut sort_push_down = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_push_down);
+ let pushed_down = pushdown_sorts(sort_push_down).unwrap().plan;
+
+ // This fallback runs at the end of EnsureRequirements, so its output must
+ // already be stable when the complete rule runs again.
+ let next_pass = optimize_and_sanity_check(Arc::clone(&pushed_down)).unwrap();
+
+ let output_children = pushed_down.children();
+ let merge_children = output_children[0].children();
+ let limit = merge_children[0]
+ .downcast_ref::<LocalLimitExec>()
+ .expect("the fetch should be carried by a local limit");
+ assert_eq!(limit.required_ordering(), &Some(ordering));
+
+ let pushed_down = plan_string(&pushed_down);
+ let next_pass = plan_string(&next_pass);
+ assert_snapshot!(pushed_down, @r"
+ OutputRequirementExec: order_by=[(a@0, asc)], dist_by=SinglePartition
+ SortPreservingMergeExec: [a@0 ASC NULLS LAST], fetch=10
+ LocalLimitExec: fetch=10
+ UnionExec
+ MockMultiPartitionExec
+ MockMultiPartitionExec
+ ");
+ assert_eq!(pushed_down, next_pass);
+}
+
+/// Verifies the complementary safe path. Before it has a fetch, a
+/// `SortPreservingMergeExec` returns one row per input row, so it may take the
+/// parent TopK fetch and forward that fetch to its input sort. Safety must be
+/// decided before adding the fetch, which changes the merge's cardinality.
+#[test]
+fn test_unfetched_spm_absorbs_and_forwards_fetch() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+ let ordering = sort_expr_on("a", 0, true, false);
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(
+ SortExec::new(ordering.clone(), source).with_preserve_partitioning(true),
+ );
+ let merge: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(ordering.clone(), sort));
+ let topk: Arc<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(ordering, merge).with_fetch(Some(10)));
+
+ let mut sort_push_down = SortPushDown::new_default(topk);
+ assign_initial_requirements(&mut sort_push_down);
+ let optimized = pushdown_sorts(sort_push_down).unwrap().plan;
+
+ let merge = optimized
+ .downcast_ref::<SortPreservingMergeExec>()
+ .expect("the merge should retain the fetch");
+ assert_eq!(merge.fetch(), Some(10));
+ let sort = merge
+ .input()
+ .downcast_ref::<SortExec>()
+ .expect("the fetch should be pushed to the input sort");
+ assert_eq!(sort.fetch(), Some(10));
+
+ // The merge retains the global fetch while the input sort retains the
+ // per-partition TopK optimization.
+ assert_snapshot!(plan_string(&optimized), @r"
+ SortPreservingMergeExec: [a@0 DESC NULLS LAST], fetch=10
+ SortExec: TopK(fetch=10), expr=[a@0 DESC NULLS LAST], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ ");
+}
+
/// Stress test: idempotency with ALL partition counts from 1 to 64
#[test]
fn test_idempotent_all_partition_counts_1_to_64() {
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
index 1f7acd8f86..19faf8ea9f 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
@@ -43,9 +43,11 @@ use datafusion_physical_plan::joins::utils::{
ColumnIndex, calculate_join_output_ordering,
};
use datafusion_physical_plan::joins::{HashJoinExec, SortMergeJoinExec};
+use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion_physical_plan::tree_node::PlanContext;
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
@@ -114,6 +116,72 @@ fn min_fetch(f1: Option<usize>, f2: Option<usize>) -> Option<usize> {
}
}
+/// Returns whether a fetch on `plan` can also be applied to each child.
+fn can_push_fetch_through(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ plan.supports_limit_pushdown()
+ && matches!(plan.cardinality_effect(), CardinalityEffect::Equal)
+}
+
+/// Returns a plan when the fast path for an already-satisfied ordering can also
+/// preserve `parent_fetch` at the current node.
+fn try_preserve_fetch_for_satisfied_plan(
+ plan: &Arc<dyn ExecutionPlan>,
+ parent_fetch: Option<usize>,
+) -> Option<Arc<dyn ExecutionPlan>> {
+ let Some(parent_fetch) = parent_fetch else {
+ return Some(Arc::clone(plan));
+ };
+ let fetch = min_fetch(plan.fetch(), Some(parent_fetch));
+
+ if plan.fetch() == fetch {
+ return Some(Arc::clone(plan));
+ }
+
+ plan.with_fetch(fetch)
+}
+
+/// Preserves a fetch above an already ordered plan while satisfying the
+/// required distribution.
+fn preserve_fetch_above_ordered_plan(
+ mut node: SortPushDown,
+ ordering: LexOrdering,
+ fetch: usize,
+ required_distribution: &Distribution,
+) -> SortPushDown {
+ // The new limit carries the fetch, so it is no longer pending on the
+ // wrapped plan.
+ node.data.fetch = None;
+
+ let input_has_multiple_partitions =
+ node.plan.output_partitioning().partition_count() > 1;
+ let input = Arc::clone(&node.plan);
+
+ let limit: Arc<dyn ExecutionPlan> = if input_has_multiple_partitions {
+ let mut limit = LocalLimitExec::new(input, fetch);
+ limit.set_required_ordering(Some(ordering.clone()));
+ Arc::new(limit)
+ } else {
+ let mut limit = GlobalLimitExec::new(input, 0, Some(fetch));
+ limit.set_required_ordering(Some(ordering.clone()));
+ Arc::new(limit)
+ };
+ let limit_node = PlanContext::new(limit, ParentRequirements::default(), vec![node]);
+
+ if input_has_multiple_partitions
+ && matches!(required_distribution, Distribution::SinglePartition)
+ {
+ let merge = SortPreservingMergeExec::new(ordering, Arc::clone(&limit_node.plan))
+ .with_fetch(Some(fetch));
+ PlanContext::new(
+ Arc::new(merge),
+ ParentRequirements::default(),
+ vec![limit_node],
+ )
+ } else {
+ limit_node
+ }
+}
+
/// Returns the stricter of two distribution requirements.
/// `SinglePartition` is the strictest.
#[expect(
@@ -244,8 +312,15 @@ fn pushdown_sorts_helper(
}
}
+ let can_push_fetch_to_children = can_push_fetch_through(&plan);
sort_push_down.plan = plan;
- if satisfy_parent {
+ let plan_with_preserved_fetch = if satisfy_parent {
+ try_preserve_fetch_for_satisfied_plan(&sort_push_down.plan, parent_fetch)
+ } else {
+ None
+ };
+ if let Some(plan) = plan_with_preserved_fetch {
+ sort_push_down.plan = plan;
// For non-sort operators which satisfy ordering:
let reqs = sort_push_down.plan.required_input_ordering();
let dists = sort_push_down
@@ -261,12 +336,19 @@ fn pushdown_sorts_helper(
} else {
parent_distribution.clone()
};
+ // A fetch retained by the current node is the pushdown boundary unless
+ // the original node can safely pass the limit through.
+ let child_fetch = if can_push_fetch_to_children {
+ parent_fetch
+ } else {
+ None
+ };
for (idx, (child, order)) in
sort_push_down.children.iter_mut().zip(reqs).enumerate()
{
child.data.ordering_requirement = order;
- child.data.fetch = min_fetch(parent_fetch, child.data.fetch);
+ child.data.fetch = min_fetch(child_fetch, child.data.fetch);
child.data.distribution_requirement = stronger_distribution(
&effective_parent_dist,
dists
@@ -306,6 +388,22 @@ fn pushdown_sorts_helper(
);
}
sort_push_down.data.ordering_requirement = None;
+ } else if satisfy_parent && let Some(fetch) = parent_fetch {
+ // Use the plan's concrete ordering when the requirement leaves sort
+ // options unspecified. If there is none, the requirement is satisfied
+ // by constants, so either direction is valid.
+ let ordering = sort_push_down
+ .plan
+ .output_ordering()
+ .cloned()
+ .unwrap_or_else(|| parent_requirement.into_single().into());
+ sort_push_down = preserve_fetch_above_ordered_plan(
+ sort_push_down,
+ ordering,
+ fetch,
+ &parent_distribution,
+ );
+ assign_initial_requirements(&mut sort_push_down);
} else {
// Can not push down requirements, add new `SortExec` (distribution-aware):
sort_push_down = add_sort_above_with_distribution(
@@ -326,14 +424,13 @@ fn pushdown_requirement_to_children(
parent_required: OrderingRequirements,
parent_fetch: Option<usize>,
) -> Result<Option<Vec<Option<OrderingRequirements>>>> {
- // If there is a limit on the parent plan we cannot push it down through operators that change the cardinality.
- // E.g. consider if LIMIT 2 is applied below a FilteExec that filters out 1/2 of the rows we'll end up with 1 row instead of 2.
- // If the LIMIT is applied after the FilterExec and the FilterExec returns > 2 rows we'll end up with 2 rows (correct).
- if parent_fetch.is_some() && !plan.supports_limit_pushdown() {
- return Ok(None);
- }
- // Note: we still need to check the cardinality effect of the plan here, because the
- // limit pushdown is not always safe, even if the plan supports it. Here's an example:
+ // A parent fetch can be pushed through a plan only if the plan explicitly
+ // supports limit pushdown and does not change cardinality. For example, if
+ // LIMIT 2 is applied below a FilterExec that removes half the rows, only one
+ // row may remain. Applied after the filter, it correctly returns two rows
+ // when at least two are available.
+ //
+ // Checking `supports_limit_pushdown()` alone is not enough. For example:
//
// UnionExec advertises `supports_limit_pushdown() == true` because it can
// forward a LIMIT k to each of its children—i.e. apply “LIMIT k” separately
@@ -347,16 +444,11 @@ fn pushdown_requirement_to_children(
// — Global LIMIT: take the first 3 rows from (A ∪ B) after merging.
// — Pushed down: take 3 from A, 3 from B, then merge → up to 6 rows!
//
- // That’s why we still block on cardinality: even though UnionExec can
+ // That’s why we also require equal cardinality: even though UnionExec can
// push a LIMIT to its children, its GreaterEqual effect means it cannot
// preserve the global TopK semantics.
- if parent_fetch.is_some() {
- match plan.cardinality_effect() {
- CardinalityEffect::Equal => {
- // safe: only true sources (e.g. CoalesceBatchesExec, ProjectionExec) pass
- }
- _ => return Ok(None),
- }
+ if parent_fetch.is_some() && !can_push_fetch_through(plan) {
+ return Ok(None);
}
let maintains_input_order = plan.maintains_input_order();
diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt
index c1e4b054c3..e17f633bc6 100644
--- a/datafusion/sqllogictest/test_files/limit.slt
+++ b/datafusion/sqllogictest/test_files/limit.slt
@@ -444,6 +444,33 @@ SELECT i FROM t1000 ORDER BY i LIMIT 3;
2
3
+# Regression test for https://github.com/apache/datafusion/issues/24927
+# The TopK fetch must not be pushed below the final DISTINCT aggregation.
+statement ok
+SET datafusion.execution.target_partitions = 2;
+
+query I
+WITH s(x) AS (
+ VALUES (0), (1), (2)
+ ORDER BY 1
+ LIMIT 3
+),
+t AS (
+ SELECT * FROM s
+ UNION ALL
+ SELECT * FROM s
+ UNION ALL
+ SELECT * FROM s
+)
+SELECT DISTINCT x FROM t ORDER BY x LIMIT 3;
+----
+0
+1
+2
+
+statement ok
+SET datafusion.execution.target_partitions = 4;
+
query I
SELECT COUNT(*) FROM (SELECT i FROM t1000 LIMIT 3);
----
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org