Skip to content

fix: retain only valid orderings after order-preserving repartition - #26747

Merged
jayzhan211 merged 6 commits into
apache:mainfrom
Kristofferjoha:fix/repartition-preserved-orderings
Nov 2, 2026
Merged

jayzhan211 merged 6 commits into
apache:mainfrom
Kristofferjoha:fix/repartition-preserved-orderings

Conversation

@Kristofferjoha

@Kristofferjoha Kristofferjoha commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #26745.
Closes #26036.

Rationale for this change

An order-preserving repartition can incorrectly retain independent input orderings that do not hold after merging partitions. This lets the planner remove a necessary sort, causing ORDER BY to return rows in the wrong order.

What changes are included in this PR?

  • Retain only the merge key ordering when repartitioning merges multiple input partitions.
  • Preserve projected cast relationships so valid cast-derived orderings can still be inferred after the merge.
  • Extend ordering checks to use equivalent expressions without assuming that equality implies a strictly order-preserving transformation.
  • Update affected expected plans.

What is the testing strategy for this PR?

  • Add the issue’s reproduction to order.slt, verifying exact row order.
  • Add safety cases in equivalence_properties.slt and monotonic_projection_test.slt.
  • Verify the existing unbounded-stream cast-ordering test continues to pass.

Are there any user-facing changes?

Only that queries affected by #26745 now return rows in the requested order.

@github-actions github-actions Bot added physical-expr Changes to the physical-expr crates sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate labels Oct 25, 2026
@Kristofferjoha

Copy link
Copy Markdown
Contributor Author

@jayzhan211 This should be ready for review. It needed a little more than following #26654, as keeping only the merge key ordering also removed valid cast-derived orderings, causing an existing unbounded-stream query to fail planning. This PR also preserves projected cast relationships so those valid orderings can be inferred again.

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Kristofferjoha thanks for picking up #26745 and for tracking down the cast-derived orderings the unbounded test needed! The repartition change looks right. Two problems with the new CastExpr rule in EquivalenceGroup::project block merging. Volatile sources get recorded as equalities, which makes ORDER BY ... LIMIT return the wrong rows. And limiting the rule to casts makes strictly order-preserving functions like to_unixtime fail planning on unbounded input.

for (target, _) in targets.iter() {
cls.push(Arc::clone(target));
}
if source.is::<CastExpr>() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This records x = CAST(r) even when the source is volatile, and then the planner picks the wrong rows. With random() the two calls are independent draws:

SET datafusion.execution.target_partitions = 1;
CREATE TABLE g AS SELECT value AS v FROM generate_series(1, 2000);
SELECT r, x FROM (
  SELECT random() AS r, CAST(random() AS FLOAT) AS x FROM g ORDER BY r LIMIT 1000
) ORDER BY x;

On main the plan is Sort[x] over TopK[r]. On this branch it becomes a single TopK(fetch=1000), expr=[x], so it returns the 1000 smallest x instead of the 1000 smallest r (max r in the output is 1.00 instead of about 0.50). Skipping volatile sources fixes it locally:

-            if source.is::<CastExpr>() {
+            if source.is::<CastExpr>() && !is_volatile(source) {

with use datafusion_physical_expr_common::physical_expr::is_volatile;.

for (target, _) in targets.iter() {
cls.push(Arc::clone(target));
}
if source.is::<CastExpr>() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Limiting this to CastExpr brings back the original planning failure for other strictly order-preserving functions. If I swap the cast in the unbounded test for to_unixtime(a) or from_unixtime(CAST(a AS BIGINT)), main plans it with a sorted window, but this branch fails with SanityCheckPlan ... Cannot execute pipeline breaking queries:

EXPLAIN SELECT a, x, SUM(c) OVER (PARTITION BY x ORDER BY b) s
FROM (SELECT to_unixtime(a) x, a, b, c FROM t_inf);

The equality holds for any deterministic expression, so recording it for every non-volatile source fixes these. The full sqllogictest suite passes with no expected-plan changes:

-            if source.is::<CastExpr>() {
+            if !source.children().is_empty() && !is_volatile(source) {

Could you add one of these queries next to the existing unbounded test?

// Equivalence Properties
let mut eq_properties = input.equivalence_properties().clone();
// If the ordering is lost, reset the ordering equivalence class:
let input_partitions = input.output_partitioning().partition_count();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional: both branches only matter when there is more than one input partition, so this could be one block, and it would keep a comment like the one in SortPreservingMergeExec:

// Fusing several input partitions keeps only the merge key ordering,
// and per-partition constants no longer hold.
if input.output_partitioning().partition_count() > 1 {
    eq_properties.clear_orderings();
    if preserve_order && let Some(ordering) = input.output_ordering() {
        eq_properties.add_ordering(ordering.clone());
    }
    eq_properties.clear_per_partition_constants();
}

@@ -1402,6 +1412,28 @@ impl EquivalenceProperties {
/// Returns an [`ExprProperties`] object containing the ordering and range
/// information for the given expression.
pub fn get_expr_properties(&self, expr: Arc<dyn PhysicalExpr>) -> ExprProperties {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional: could the doc comment mention that this now falls back to the sort properties of an equivalent expression? A short comment on the new block in class.rs saying why the equality is recorded would also help: it keeps orderings derived through strict functions after a merge resets the orderings.

@jayzhan211
jayzhan211 marked this pull request as draft October 26, 2026 02:27
@github-actions github-actions Bot added the optimizer Optimizer rules label Oct 26, 2026
@Kristofferjoha

Copy link
Copy Markdown
Contributor Author

@jayzhan211 Thank you for the review as always. Addressed the comments. While testing, I found another wrong-result case exposed by the broader equality rule.

SET datafusion.execution.target_partitions = 1;

CREATE TABLE lossy_topk AS VALUES (16777217), (16777216);

SELECT r, x FROM (
  SELECT column1 AS r, CAST(column1 AS FLOAT) AS x
  FROM lossy_topk
  ORDER BY r
  LIMIT 1
) ORDER BY x;

This should return r = 16777216, but with the broader equality rule it returns r = 16777217. Both integers cast to the same float value, and sort pushdown replaces the inner TopK[r] with TopK[x], changing which row survives the limit.

I added a small fix to retain the inner TopK’s selection keys when they already satisfy the outer ordering, plus an exact-result regression test.

@Kristofferjoha
Kristofferjoha marked this pull request as ready for review October 31, 2026 19:10

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Kristofferjoha thanks for the updates, and nice catch on the lossy TopK case! Everything from the last round is fixed. One problem remains, and it comes from the broader rule I suggested last time. Details inline.

}
// Retain deterministic expressions as equalities so orderings derived
// through strict functions can be inferred after a merge clears orderings.
if !source.children().is_empty() && !is_volatile(source) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The broader rule I suggested last round has a problem, sorry about that: the new equality survives an outer join's NULL padding. This is #26036, which so far needed a filter like WHERE b = coalesce(a, 0). With this PR a plain projection is enough. On a padded row x is NULL but CASE WHEN a IS NULL THEN 0 ELSE a END is 0, so on this branch the outer sort below is dropped and the query returns 1, 2, 0. Main returns 0, 1, 2:

SET datafusion.execution.target_partitions = 1;
CREATE TABLE padded_l(k INT) AS VALUES (1), (2), (3);
CREATE TABLE padded_r(k INT, a INT) AS VALUES (1, 1), (2, 2);
SELECT y FROM (
  SELECT CASE WHEN a IS NULL THEN 0 ELSE a END AS y FROM (
    SELECT rr.x, rr.a
    FROM (SELECT k, a, CASE WHEN a IS NULL THEN 0 ELSE a END AS x FROM padded_r) rr
    RIGHT JOIN padded_l ON padded_l.k = rr.k
    ORDER BY rr.x LIMIT 100
  )
) ORDER BY y;

Keeping only the column members on the padded side in EquivalenceGroup::join fixes both this query and the #26036 repro. With it and the query above added to equivalence_properties.slt (expected 0, 1, 2), the full sqllogictest suite passes with no expected-plan changes. Could you include it here and add Closes #26036?

-use crate::expressions::Literal;
+use crate::expressions::{Column, Literal};
@@ impl EquivalenceClass {
+    /// Keeps what still holds on a join side whose unmatched rows are padded
+    /// with NULLs: padded columns stay equal to each other, but constants and
+    /// computed members do not (`coalesce(a, 0)` is 0 where `a` is NULL).
+    fn null_pad(&mut self) {
+        self.constant = None;
+        self.exprs.retain(|expr| expr.is::<Column>());
+    }
@@ pub fn join(
                             if matches!(join_type, JoinType::Right | JoinType::Full) {
-                                cls.constant = None;
+                                cls.null_pad();
                             }
@@
                                     ) {
-                                        cls.constant = None;
+                                        cls.null_pad();
                                     }

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Outer sort should be retained now, added Closes #26306 to description.

@jayzhan211

Copy link
Copy Markdown
Contributor

Thanks @Kristofferjoha !

@jayzhan211
jayzhan211 added this pull request to the merge queue Nov 2, 2026
Merged via the queue into apache:main with commit 99e723d Nov 2, 2026
44 checks passed
geoffreyclaude pushed a commit to geoffreyclaude/datafusion that referenced this pull request Nov 12, 2026
…rdering (apache#27191)

## Which issue does this PR close?

- Closes apache#26762.

## Rationale for this change

Queries that planned in earlier releases fail with `SanityCheckPlan`
when a window's input ordering starts with a per-partition constant or
is proven through an equality such as `a = CAST(c AS DATE)`. Across
7,440 generated query runs, `main` has 62 failures on queries that plan
on 55.2.0, and this PR has none, with the same rows wherever both plan.

There are three causes. In each one, window planning and
`SanityCheckPlan` reach different answers about whether the input is
sorted:

| Introduced by | Example | Why they disagree |
|---|---|---|
| apache#23528 (in 55.0.0) | The apache#26762 query | Normalized orderings keep
per-partition constants, but `add_constants` drops them, so window
planning sees `[total]` and `SanityCheckPlan` sees `[k, total]` |
| apache#25922 (on `branch-56`) | A window on `c` over a TopK on `(k, a, c)`,
with `k` a per-partition constant and `a = CAST(c AS DATE)` |
`add_constants` discovers that `[a, c]` implies `[c]`, but an ordering
added after the equality doesn't |
| apache#26747 (on `branch-56`) | A window on `(k, c)` over a TopK on `(k, a,
c)`, where `a` isn't selected | The sort is removed below a projection
because its input has the ordering, and the projection can't prove it |

This PR replaces apache#27140 and keeps its commit.

## What changes are included in this PR?

Start reading at `add_orderings`. It is the decision I most want
reviewed, because it makes the optimizer claim more orderings, and a
wrong claim means wrong results.

| Where | Change | Fixes |
|---|---|---|
| `SortKeyWalk::new` | An ordering's head skips leading fixed
expressions before the first key is pinned | apache#23528 |
| `add_orderings` | Discover the orderings a new ordering implies
through known equalities, as `add_constants` and `add_equal_conditions`
already do, for any key and past the constants an ordering keeps: `[k,
a, c]` with `a = CAST(c AS DATE)` implies `[k, c]` | apache#25922, apache#26747 |
| `handle_projection_pushdown` | If the input has the ordering but the
projection can't prove it, keep the sort above the projection | The
failures the first row would otherwise add, when a projection drops a
leading per-partition constant |
| `RepartitionExec` | Clear per-partition constants before adding the
merge ordering, as `SortPreservingMergeExec` does | Otherwise the second
row keeps `[c]` from `[k, a, c]` after a merge, where it holds only
within each partition |
| `SortKeyWalk::is_fixed` | Return early for a column that isn't pinned
or constant | The first row made single-key checks about 30% slower |

**Who is affected.** Plans with an ordering that starts with a
per-partition constant, or that is proven through an equality or a
column a projection drops. It can't stop halfway.

**No API change.** Discovery in `add_orderings` can fail only on an
expression that doesn't match the schema. That implied ordering is then
left out, which can cost a sort but not change results.

**Worst case.** 4 of the 7,440 generated runs plan a sort that `main`
doesn't: a `GROUP BY k, c` over a TopK on `(a, c)` with `a = CAST(c AS
DATE)`, at 4 partitions. Knowing `[c]`, the aggregate streams on `c` and
sorts on `c` after its hash repartition. `main` picks the same aggregate
when the ordering is known before the equality. 392 runs have fewer
sorts.

## What is the testing strategy for this PR?

- `window.slt` covers each cause, including a merge of partitions above
the sort, and checks the plans. Removing any one of the first three
changes, limiting discovery to the first key, or not skipping constants
makes tests fail, in `window.slt` or in the unit tests for the walk and
for implied orderings. A `RepartitionExec` unit test covers the merge.
- All SLT files pass with no existing plan changing, as do the
`physical_optimizer` integration tests and the equivalence fuzz tests.
- A sweep of functions that look order-preserving but aren't found no
wrong ordering. Details below.

<details>
<summary>Generated query runs, benchmarks, rejected alternatives, and
queries that still fail</summary>

**Generated runs.** Each sweep runs at `target_partitions` 1 and 4. A
regression plans on the reference build and fails on the other build.

| Sweep | Runs | `main`, vs 55.2.0 | `main`, vs before apache#23528 | This PR,
vs both |
|---|---|---|---|---|
| Per-partition constants: 6 inputs, 3 orderings, 5 projections, up to
12 consumers | 2,016 | 0 | 4 | 0 |
| Orderings proven through `a = CAST(c AS DATE)` | 3,120 | 56 | 56 | 0 |
| An operator between the ordering and the consumer: joins of each type
on either side, filter, aggregate, union | 1,584 | 6 | 36 | 0 |
| `head = f(arg)` for `CAST(x AS VARCHAR)`, `ABS`, negation, modulo,
`upper` and order-preserving functions, in all 4 sort directions | 720 |
0 | 0 | 0, and no wrong ordering |

The results are the same with this patch on `branch-56`, where all SLT
files pass as well. Compared with `main` on the runs that plan on both,
the rows are the same, 4 runs have more sorts (see the worst case above)
and 392 have fewer.

**Benchmarks.** `equivalence_properties`, best of 10 alternating runs on
a busy laptop, so approximate: single-key checks are 6-12% slower (about
20-40 ns), and window-heavy orderings are 24-33% faster.

| Benchmark | `main` | This PR |
|---|---|---|
| `ordering_satisfy/1_key/8` | 359 ns | 380 ns |
| `ordering_satisfy_requirement/1_key/8` | 307 ns | 329 ns |
| `find_longest_permutation/1_expr/32` | 306 ns | 344 ns |
| `window/4` | 1,543 ns | 1,181 ns |
| `window/32` | 7,319 ns | 4,889 ns |

**Rejected.**

- Having `add_constants` keep per-partition constants fails
`test_eliminate_redundant_monotonic_sorts`.
- Labelling `lower(g)` as a uniform constant isn't enough, because a
constant from an aggregate without `GROUP BY` really is per-partition.
- Having the projection leave a per-partition constant it drops out of
its orderings, instead of the check in `handle_projection_pushdown`,
makes 2 generated queries fail `SanityCheckPlan`: a window over a `GROUP
BY` whose input ordering leads with such a constant.
- Applying the projection check to every node in `pushdown_sorts_helper`
only fixes one more query, and that query already fails on every
release, so I kept the narrower check.

**Still failing, not regressions.** A window over a `UNION ALL` of two
TopKs that have `v AS w` fails on every build I tried, including before
apache#23528:

```sql
SELECT *, ROW_NUMBER() OVER (ORDER BY w) AS r FROM (
  SELECT k, v, w FROM (SELECT k, v, v AS w FROM (SELECT lower(g) AS k, v FROM t5 WHERE g = 'a') ORDER BY k, v LIMIT 30)
  UNION ALL
  SELECT k, v, w FROM (SELECT k, v, v AS w FROM (SELECT lower(g) AS k, v FROM t5 WHERE g = 'a') ORDER BY k, v LIMIT 30)
) ORDER BY r;
```

</details>

## Are there any user-facing changes?

Affected queries plan again. This is a candidate for `branch-56`, where
it applies without conflicts.

---------

Co-authored-by: Andrew Lamb <[email protected]>
Co-authored-by: Claude Fable 5.1 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

optimizer Optimizer rules physical-expr Changes to the physical-expr crates physical-plan Changes to the physical-plan crate sqllogictest SQL Logic Tests (.slt) v56.0.0

Projects

None yet

2 participants