cning112 opened a new pull request, #26455:
URL: https://github.com/apache/datafusion/pull/26455
## Which issue does this PR close?
- Closes #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, `sort_ordering_above`, which
still removes them, with one exception. When the new `SortExec` has a `fetch`
and its input has more than one partition, a constant that differs between
partitions (`AcrossPartitions::Heterogeneous`) stays in the key. The shared
threshold then compares the full key, which is comparable across partitions.
Nothing else changes:
- `Uniform` constants are still removed.
- Sorts without a `fetch` still remove `Heterogeneous` constants. A full
per-partition sort has no shared threshold, and the merge above it sorts on the
full key.
- With a single input partition the sort is not partitioned, so nothing is
shared.
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.
#23528 fixed the same `Uniform`/`Heterogeneous` confusion in
`EquivalenceGroup`'s sort normalization. This PR deals only with the fetch case
in these two helpers.
Open PRs #25408 and #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`.
## What is the testing strategy for this PR?
New cases at the end of `datafusion/sqllogictest/test_files/topk.slt`, on a
`VALUES` table:
- the query above with an `EXPLAIN` showing the TopK key, at
`target_partitions = 1` and `4`;
- the same query without the filter;
- `LIMIT 3` and `ORDER BY c0 DESC, c1` as controls.
On `main` without the fix, four of them fail: the `EXPLAIN`, `LIMIT 2` at
both partition counts, and the no-filter query. The `LIMIT 3` and `DESC`
controls pass both before and after.
No existing `.slt` plan changes. `--complete` over the whole suite rewrote
nothing but trailing newlines in three unrelated files, and I reverted those.
## 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 #22991 for these plans.
My reasoning is that the exit stays correct because it compares on the full key
prefix, which now includes `c0`. The `DESC` case exercises the opposite
partition order, but I have not written a test that forces either partition to
finish first.
- **Merge above the sort.** `add_sort_above_with_distribution` wraps the
sort in a `SortPreservingMergeExec` with the same shortened ordering when the
parent requires a single partition. Without a fetch, that merge could also lose
a `Heterogeneous` key. I could not build a query that reaches it, and this PR
does not change it.
- **Other places that attach a fetch.** I assume no other rule attaches a
fetch to a partitioned `SortExec` whose key was already shortened. In the
queries I tried, the fetch is already set when `EnsureRequirements` builds the
sort: plain `LIMIT`, `LIMIT ... OFFSET`, and a filter over the union.
`LimitPushdown` left those plans unchanged. I have not audited every caller of
`SortExec::with_fetch`.
- **The no-filter variant.** Its result on `main` depends on how the
partitions run: wrong in the sqllogictest, right over a single-batch `MemTable`
in my runs. I did not dig into why.
- **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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
