NoahKusaba opened a new issue, #2854:
URL: https://github.com/apache/datafusion-ballista/issues/2854
**Describe the bug**
With the adaptive planner (the default), a `UNION ALL` whose branches each
end in a node that merges its input into one partition, such as the `SortExec`
an `ORDER BY` puts on each branch, returns some rows several times when the
union's stage is split between tasks.
When a stage is split, each task gets the plan restricted to its own
partitions, and a branch of a `UnionExec` the task doesn't own is restricted to
no partitions. A branch that merges its input into one partition still reports
one partition afterwards, and the nodes below the merge read their input whole.
The executor runs every partition of the task's plan, so each task runs that
branch in full, and its rows come back once per task.
The static planner already handles this: `can_stay_inline`
(`ballista/scheduler/src/planner.rs`) gives such a branch a stage of its own,
so it reaches the union as a shuffle reader that restricts to nothing cleanly
(the fix for #2184). The adaptive planner has no equivalent.
Whether a stage is split depends on free task slots, so this shows on a busy
cluster, or always with `ballista.scheduler.max_partitions_per_task = 1`.
**To Reproduce**
Standalone Ballista, a one-file Parquet table `p` with rows `(1, 'a'), (2,
'b')`, and `ballista.scheduler.max_partitions_per_task = 1`:
| Query | Rows | Expected |
|---|---|---|
| `SELECT id, name FROM p UNION ALL SELECT id, name FROM p` | 4 | 4 |
| `SELECT id, name FROM p UNION ALL SELECT id, name FROM p ORDER BY id` | 8
| 4 |
Other branches that merge their input into one partition do the same. Each
was checked against DataFusion on the adaptive planner with
`max_partitions_per_task = 1`:
| Branch shape | Example | Result |
|---|---|---|
| aggregate without `GROUP BY` | `SELECT count(*) FROM test WHERE id > 100
UNION ALL SELECT count(*) FROM test` | 3 rows, should be 2 |
| `LIMIT` (`CoalescePartitionsExec`) | `SELECT count(*) FROM (SELECT o_id
FROM (SELECT o_id FROM orders LIMIT 10) UNION ALL SELECT c_id FROM customers)`
| over 110, should be 110 |
| `ORDER BY … LIMIT` (`SortPreservingMergeExec`) | `SELECT * FROM (SELECT
o_id FROM orders ORDER BY o_id LIMIT 3) UNION ALL SELECT c_id FROM customers
WHERE c_id < 2` | 7 rows, should be 5 |
(`test` is `alltypes_plain.parquet`; `orders` is 1000 rows over four CSV
files and `customers` 100 rows in one.)
The same happens with an Iceberg table (andygrove's repro on #2638).
**Expected behavior**
4 rows, as without the split.
**Additional context**
- #2852 (open) makes the scheduler refuse such a task ("cannot restrict a
`UnionExec` branch to partitions []: its `SortExec` keeps a partition count of
1…"), so the query fails instead of returning wrong rows. This issue is for
making it succeed. Until then, `ballista.planner.adaptive.enabled = false`
avoids it.
- The scheduler can't fix it per task: replacing the whole branch with an
empty node changes the task plan's node count, which `update_task_metrics`
rejects, and emptying only the branch's leaves still leaves a global
aggregate's `count = 0` row.
- Proposed fix: have the adaptive planner give a union branch that can't be
restricted to no partitions a stage of its own, as the static planner does, by
sharing `can_stay_inline`'s test (a trial restriction to `[]`, plus the
broadcast check) with `DistributedExchangeRule`
(`ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs`).
- Related: #2184 (the static planner's fix), #2679 / #2721 (leaves the
scheduler can't restrict).
--
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]
