cning112 opened a new issue, #26454:
URL: https://github.com/apache/datafusion/issues/26454
### Describe the bug
An `ORDER BY ... LIMIT` over a `UNION ALL` returns wrong rows when a leading
sort key is constant within each arm but has a different value in each arm.
Reproduced on `main` at 5f2158269 (2026-10-16) with the statements below.
The same query over a one-partition `MemTable` returns the same wrong rows and
plan on `main` at 5f2158269 and cdc2fc041, and on the 55.2.0 release.
### To Reproduce
As sqllogictest statements:
```sql
CREATE TABLE t(id BIGINT, g BIGINT) AS VALUES (0, 0), (1, 0), (2, 1), (3,
1), (4, 1);
SET datafusion.execution.target_partitions = 1;
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;
```
This returns `(1, 0), (5, 0)`. The expected result is `(1, 0), (1, 1)`: the
second arm has two rows with `c0 = 1`, and they sort before every row with `c0
= 5`. The same query returns the wrong rows with `target_partitions = 4`.
The physical plan:
```
SortPreservingMergeExec: [c0@0 ASC NULLS LAST, c1@1 ASC NULLS LAST], fetch=2
SortExec: TopK(fetch=2), expr=[c1@1 ASC NULLS LAST],
preserve_partitioning=[true]
UnionExec
ProjectionExec: expr=[5 as c0, id@0 as c1]
DataSourceExec: partitions=1, partition_sizes=[1]
ProjectionExec: expr=[1 as c0, id@0 as c1]
FilterExec: g@1 = 0, projection=[id@0]
DataSourceExec: partitions=1, partition_sizes=[1]
```
The TopK key is `[c1]`. `c0` was dropped from it.
`LIMIT 3` on the same query returns the right rows. Without `WHERE g = 0`
the result depends on how the partitions run: the sqllogictest above returns
`(1, 0), (5, 0)`, while the same query over a single-batch `MemTable` returned
the right rows in my runs.
### Expected behavior
`(1, 0), (1, 1)`.
### Additional context
### Analysis
Two pieces combine:
1. `EnsureRequirements` builds the partitioned `SortExec` with the fetch
already attached. `EXPLAIN VERBOSE` shows the shortened `SortExec:
TopK(fetch=2), expr=[c1@1 ASC NULLS LAST]` in the plan right after
`EnsureRequirements`. With a print in the helpers, it is built by
`add_sort_above` (from `parallelize_sorts`, through
`add_sort_above_with_check`), and built again by
`add_sort_above_with_distribution` (from `pushdown_sorts`). Both helpers, in
`datafusion/physical-optimizer/src/utils.rs`, drop every sort key for which
`is_expr_constant(..)` is `Some`. They make no distinction between
`AcrossPartitions::Uniform` and `AcrossPartitions::Heterogeneous`. After the
`UnionExec`, `c0` is `Heterogeneous`: 5 in one partition and 1 in the other.
For a per-partition sort feeding a `SortPreservingMergeExec` on the full key,
dropping it is fine.
2. All partitions of that `SortExec` share one TopK threshold.
`SortExec::create_filter` creates one `TopKDynamicFilters` for the operator,
and `TopK::insert_batch` filters its own input with it. Partition 0 (`c0 = 5`)
fills its heap with `c1 ∈ {0, 1}` and publishes `c1 < 1`. Partition 1 (`c0 =
1`) then drops its `(1, 1)` row. A threshold on the shortened key cannot be
compared across partitions whose dropped constant differs.
Setting `datafusion.optimizer.enable_topk_dynamic_filter_pushdown = false`
and `datafusion.optimizer.enable_dynamic_filter_pushdown = false` does not
change the result, because `TopK::insert_batch` applies the shared filter
regardless of those settings.
Possible fixes:
- Keep `Heterogeneous` constants in the key of a partitioned `SortExec` that
has a fetch, so the shared threshold compares the full key. I have a small
patch for this and can open a PR.
- Or stop sharing the TopK threshold across partitions when the input has a
`Heterogeneous` constant. That needs a TopK change and gives up the shared
filter for those plans.
### Related
- #23528 fixed the same `Uniform`/`Heterogeneous` confusion in
`EquivalenceGroup`'s sort normalization. The two optimizer helpers were not
changed.
- #24927 (fixed by #24932) was another wrong-rows bug in the same plan
shape, a `SortPreservingMergeExec` over a partitioned TopK. The cause was
different: a fetch was pushed below a final aggregation.
- Open PRs #25408 and #25065 change the callers of these helpers
(`sort_pushdown.rs` and `enforce_sorting/mod.rs`), but not the helpers
themselves.
### Additional context
Found while testing rustq, a SQL front end on DataFusion. Our workaround is
a physical optimizer rule that removes the fetch from a partitioned `SortExec`
whose input has a heterogeneous constant.
### Willingness to contribute
I plan to contribute a fix for this bug independently
--
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]
