Distributed execution returns inflated results and crashes on UNION queries at target_partitions=1

[andygrove]

Describe the bug

At target_partitions=1, tasks of a UnionExec stage each execute every union branch instead of only the branch they own. That shows up two ways: silently inflated results, and a hard failure.

Surfaced by extending the TPC-DS SF1 correctness gate to run at 1 partition as well as 16 (#2182).

Confirmed cases (SF1, --partitions 1, prefer_hash_join=false, static planner)

Query	Single-process DataFusion	Ballista (distributed)	Ratio
q5	73724.99	294899.96	exactly 4x
q80	11343.60	22687.20	exactly 2x

q49 does not return a wrong answer, it fails outright:

Job failed due to stage 7 failed: Task failed due to runtime execution error:
Internal("FileStreamBuilder invalid partition index: 0")

q5, q80 and q49 are exactly the three queries in the affected set that contain UNION. All three pass at --partitions 16.

Root cause

Confirmed by instrumenting restrict_plan_to_partitions and capturing the real plans, not by inspection.

Restricting q49's 3-partition union stage to [0] returns a plan that still has 3 partitions. The executor (drive_shuffle_writer_stage) runs 0..partition_count, so every task runs all three branches.

Why the restriction fails to shrink it: a union branch that collapses internally reports one output partition no matter what its leaves do. In these queries the collapse is

• CoalescePartitionsExec (q5's store and web branches),
• a CollectLeft HashJoinExec (q5's catalog branch),
• SortExec: preserve_partitioning=[false] (q49's rank() over branches).

restrict allocates the parent's partitions across children, so an unowned branch is restricted to [] — but under_collect scoping then deliberately leaves every leaf below the collapse reading its full upstream, and broadcast ShuffleReaderExecs are exempted too. The branch keeps reporting a partition and keeps producing all of its data.

Hence N tasks x N branches: N-fold duplication (q5, q80), or a task executing a branch whose scan was emptied (q49).

Why this is not simply a gap in #2070

#2070 fixed a different failure in the same area (a scan reading its whole table because the stage partition index was applied directly to every scan) and verified partition counts 4, 5, 6, 8, 16, 20, 32 — never 1. This bug is adjacent but distinct: the per-child index arithmetic is correct; the problem is that an unowned branch cannot be reduced to zero partitions.

Why it cannot be fixed inside the restriction

Three constraints that are mutually incompatible for this shape. Two of them were hit as real runtime failures while attempting fixes:

1. The restricted plan must expose exactly indices.len() partitions. compute_global_output_partition_ids returns PassThrough(global_input_partition_ids) when the child chain branches, and the executor maps its local output slots onto those global ids.
2. Operator count must be identical across a stage's tasks. ExecutionStage::update_task_metrics matches metrics positionally and rejects a length mismatch, so a rewrite cannot add or remove plan nodes. Replacing an idle branch with a 0-partition stub produces task metrics array size 9 does not equal with the stage metrics array size 12 and the job then stalls.
3. Operators require at least one input partition. Draining an idle branch's leaves to zero partitions instead produces CoalescePartitionsExec requires at least one input partition (q5) and Invalid HashJoinExec, the output partition count of the left child must be 1 in CollectLeft mode (q49).

So an unowned collapsing branch cannot be made 0-partition by restriction (1), cannot be structurally replaced (2), and cannot have its leaves zeroed (3).

Proposed fix

Relax constraint 1: give the executor the task's assigned partition indices and have it execute exactly those, rather than every partition of the plan it is handed. The union then needs no structural rewrite at all, operator count stays identical across tasks, and no operator is handed zero partitions. Leaf restriction is still needed so a scan does not drain DataFusion's shared work queue.

To Reproduce

Generate TPC-DS SF1 Parquet with tpcgen-cli, start a scheduler and executor, then:

tpcds --host localhost --port 50050 --path <data> --partitions 1 --verify \
  --query 5 -c datafusion.optimizer.prefer_hash_join=false
# -> Query 5 VERIFY MISMATCH: expected `73724.99`, got `294899.96`

tpcds --host localhost --port 50050 --path <data> --partitions 1 --verify \
  --query 49 -c datafusion.optimizer.prefer_hash_join=false
# -> Internal("FileStreamBuilder invalid partition index: 0")

Both pass with --partitions 16.

Expected behavior

Results at target_partitions=1 should match single-process DataFusion, as they do at every other partition count.