Fix ordering for UNION ALL over heterogeneous constants - #23528

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes apache#123` indicates that this PR will close issue apache#123.
-->

- Closes #.

## Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->

`add_sort_above` and `add_sort_above_with_distribution` in
physical-optimizer prune every expression that
`EquivalenceProperties::is_expr_constant` reports as constant.
This treats `AcrossPartitions::Heterogeneous` the same as
`AcrossPartitions::Uniform`, which is incorrect.

Check an example:

```sql
SELECT 2 AS sample, c FROM t
UNION ALL
SELECT 1 AS sample, c FROM t
ORDER BY sample, c
```

The union reports sample as a `Heterogeneous` constant and so it gets
pruned, which breaks `ORDER BY`.

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

Helper that activates pruning only for `AcrossPartitions::Uniform`, use
it at both prune sites (`add_sort_above` and
`add_sort_above_with_distribution`) so `Heterogeneous` constants are
retained.

## Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

Yes, a new unit-test added.

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

No API changes, just an internal logic bug-fix.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>

…che#26455)

## Which issue does this PR close?

- Closes apache#26454.

## Rationale for this change

`ORDER BY ... LIMIT` over a `UNION ALL` returns wrong rows when a
leading sort key is a constant with a different value in each arm:

```sql
SELECT 5 AS c0, id AS c1 FROM t
UNION ALL
SELECT 1 AS c0, id AS c1 FROM t WHERE g = 0
ORDER BY c0, c1 LIMIT 2
-- returns (1, 0), (5, 0); expected (1, 0), (1, 1)
```

The optimizer plans a partitioned TopK under a `SortPreservingMergeExec`
and drops `c0` from the TopK key, because `c0` is constant within each
partition. All partitions of one `SortExec` share a single TopK dynamic
filter threshold. The `c0 = 5` partition publishes a threshold on `c1`
alone, and the `c0 = 1` partition then discards a row it needs. The
issue has the full walk-through.

## What changes are included in this PR?

`add_sort_above` and `add_sort_above_with_distribution` in
`datafusion/physical-optimizer/src/utils.rs` both remove sort keys that
are constant in the input. They now share one helper,
`ordering_without_constants(plan, sort_requirements,
keep_heterogeneous)`, which still removes them, except that a constant
whose value differs between partitions
(`AcrossPartitions::Heterogeneous`) is kept when `keep_heterogeneous` is
set. Each caller sets it when an operator compares rows from different
partitions on the new ordering:

- **Partitioned TopK** (both helpers): the new `SortExec` has a `fetch`
and its input has more than one partition. The shared threshold then
compares the full key, which is comparable across partitions.
- **Merge above the sort** (`add_sort_above_with_distribution` only):
when the parent requires a single partition, the helper wraps the sort
in a `SortPreservingMergeExec` built from the same ordering. That merge
compares rows across partitions even without a fetch, and since it
clears per-partition constants from its output, a shortened key would no
longer guarantee the required ordering. The constant is now kept there
too.

The helper preserves these existing behaviors:

- `Uniform` constants are still removed.
- A partitioned sort without a `fetch` and without an added merge still
removes `Heterogeneous` constants. Each partition sorts on its own, and
any merge above it sorts on the full key.
- With a single input partition nothing is compared across partitions.

With the change, the plan for the query above is:

```
SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
  SortExec: TopK(fetch=2), expr=[c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], preserve_partitioning=[true], sort_prefix=[c0@0 ASC NULLS LAST]
```

The longer key gives the TopK a `sort_prefix` of `[c0]`, because each
partition's input is trivially ordered by its constant `c0`.

Why the fix lives here: `EXPLAIN VERBOSE` shows the shortened `SortExec:
TopK(fetch=2), expr=[c1@1 ...]` already in the plan right after
`EnsureRequirements`, with the fetch attached. No later rule changes it.
A temporary print in the two helpers showed `parallelize_sorts` building
it through `add_sort_above_with_check` → `add_sort_above`, and
`pushdown_sorts` building it again through
`add_sort_above_with_distribution`. Both calls had `fetch = Some(2)`,
two input partitions, and `c0` reported as `Heterogeneous`. So both
helpers need the change, and fixing only one would leave the bug.

apache#23528 fixed the same `Uniform`/`Heterogeneous` confusion in
`EquivalenceGroup`'s sort normalization. This PR applies it to the TopK
key and to the merge that these two helpers build.

Open PRs apache#25408 and apache#25065 change the callers of these helpers
(`enforce_sorting/sort_pushdown.rs` and `enforce_sorting/mod.rs`) but
not `utils.rs`. This PR should not conflict with them textually, but
whichever merges second should re-run `topk.slt`.

Two separate validation prerequisites are included: three narrowly
scoped `mutable_key_type` Clippy annotations matching existing optimizer
conventions, and lockfile patch updates for `rustls-webpki` (0.103.16)
and `webpki-root-certs` (1.0.10). The exact all-features Clippy command
failed on the unchanged base, and the security audit exposed
RUSTSEC-2026-0382 / RUSTSEC-2026-0383 in the existing lockfile. Neither
prerequisite changes the TopK implementation.

## What is the testing strategy for this PR?

- **`topk.slt`**, on `VALUES` tables, each with an `EXPLAIN` that pins
the TopK key and a result query:
  - the query above, and the same query without the filter;
  - `ORDER BY c0 DESC, c1`;
- `ORDER BY c1, c0`, where the constant is not the leading key and only
breaks ties on `c1`;
- `LIMIT ... OFFSET`, where the TopK needs three rows from the `c0 = 1`
arm;
- a three-arm union, so two arms can publish a threshold that cuts the
third;
- controls: `LIMIT 3`, a `Uniform` constant (end-to-end control), and no
`fetch` (sort pushed into the arms).

Whether a result query fails without the fix depends on which arm
publishes its threshold first (the `DESC` query never failed in 8 runs),
so the `EXPLAIN`s are what catch a regression reliably.
- **Unit tests in `utils.rs`**, one per branch: a partitioned TopK keeps
a `Heterogeneous` constant; the added merge keeps it with and without a
fetch; a partitioned sort without a fetch or merge drops it; a `Uniform`
constant is always dropped.
- **Unit tests in `topk/mod.rs`**: two TopKs share one
`TopKDynamicFilters`, with key `[c0, c1]` and `sort_prefix` `[c0]`. Both
feeding orders return the right top 2. A second test feeds several
batches per partition and checks the early-exit state: the better
partition keeps reading until it has `k` rows, and only then does the
worse one finish.

Ablation, with all tests in place:

- Dropping `Heterogeneous` constants again in
`ordering_without_constants` fails both `keep` unit tests in `utils.rs`
and 10 `topk.slt` assertions, including every new `EXPLAIN`.
- Removing only the merge part (`|| add_merge`) fails
`merge_keeps_a_constant_that_differs_between_partitions` on its `fetch =
None` case. No `.slt` test fails, see below.

No existing `.slt` plan changes.

Local validation on the final tree: formatting, `cargo clippy
--all-targets --all-features -- -D warnings`, and `uv run
./dev/rust_lint.sh` pass. The required extended workspace test command
passes with 12,733 Rust tests (8 ignored) and all 535 SQL files.
Independent mutations also confirm the new TopK tests fail when the
leading filter key is omitted, and the multi-batch test fails when the
shared-prefix comparison is reversed. Benchmark datasets were
unavailable; no performance claim is made.

## Are there any user-facing changes?

The wrong results are fixed. Plans that hit this case now show the
constant in the TopK `expr`. There are no API changes.

## Unknowns / assumptions

- **Shared prefix threshold.** The new key gives the TopK a
`sort_prefix` of `[c0]`. That enables the shared-prefix early exit from
apache#22991 for these plans. It stays correct because it compares on the full
key prefix, which now includes `c0`, and a threshold is only published
once the publishing partition has `k` rows. The multi-batch TopK unit
test exercises this directly; no `.slt` test forces a particular
partition to finish first.
- **No tested SQL reached the no-fetch merge branch.** I instrumented
`add_sort_above_with_distribution` for `fetch = None` with a
single-partition requirement and more than one input partition. It
recorded no hits across the whole sqllogictest suite, for candidate
queries (un-partitioned windows over union, filter, join, aggregate and
distinct), or with `optimizer.repartition_sorts = false` /
`optimizer.enable_sort_pushdown = false`. Distribution enforcement and
`parallelize_sorts` build the merge with the full ordering before
`pushdown_sorts` runs. That is an observation, not a proof, so that
branch is covered by the unit test only.
- **Fetch attached after key simplification (accepted risk;
follow-up).** The helper decides whether to retain heterogeneous
constants when it builds the sort. A partitioned sort without a fetch or
helper-added merge may omit `c0`; if `LimitPushdown` or
`LimitPushPastWindows` subsequently calls `SortExec::with_fetch`, that
method enables shared TopK filtering without restoring `c0`. This PR
does not guard that transition. No SQL query tried so far reaches the
entire sequence under the current rule order: the tested shapes already
have a fetch at construction, push sorts into individual union arms,
retain the full key below a merge, hash-repartition first, or eliminate
the inner sort. That evidence does not establish unreachability. The
proposed follow-up is a hand-built plan regression plus an executor
guard that isolates or disables unsafe shared filtering, including
dynamic-filter pushdown. The reproduction is proposed, not yet executed.
`preserve_fetch_above_ordered_plan` remains an additional path to audit.
- **Alternative fix.** The other option is to stop sharing the TopK
threshold across partitions when the input has a `Heterogeneous`
constant. That would be a larger change inside TopK and would give up
the shared filter. I went with the smaller planner change, but I am
happy to switch if maintainers prefer it.

---------

Co-authored-by: Claude Opus 5.5 <[email protected]>
Co-authored-by: Jay Zhan <[email protected]>

…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]>