Skip the data shuffle for UNION ALL in the multi-stage engine

[yashmayya]

UNION ALL is a pure concatenation, so any row-to-worker mapping produces the same result and no redistribution is ever required. The V2 physical optimizer has always known this — TraitAssignment#assignSetOp returns early for Union.all ("UNION ALL means we can return duplicates, so no trait required"). The V1 (default) planner did not: PinotSetOpExchangeNodeInsertRule unconditionally hash-shuffled every set-op input on the full output row, UNION ALL included.

This brings V1 in line. A UNION ALL input now gets a local exchange — SINGLETON, which is how Pinot already spells "local exchange" — carrying the projected columns as keys. The union stage inherits its inputs' worker assignment and each sender hands its rows to the worker on its own server: no shuffle, no hashing, and no network hop when sender and receiver land together.

Opt out per query with /*+ setOpOptions(is_colocated_by_set_op_keys='false') */, which restores the full-row shuffle.

Why the keys are attached

The keys are unused while the exchange stays local. They are there so the mailbox layer can promote the exchange to a real HASH_DISTRIBUTED shuffle when the branches cannot share a worker assignment — the same idiom PinotJoinExchangeNodeInsertRule already uses for distribution_type='local':

// NOTE: We use SINGLETON to represent local distribution. Add keys to the exchange because we might want to switch it to HASH distribution to increase parallelism.

That also keeps the change from having to distinguish itself from anything else. SINGLETON is shared with two features whose contracts are much stricter — the colocated dynamic-broadcast semi-join build side (every receiver needs the whole build side; the non-colocated variant broadcasts for exactly that reason) and distribution_type='local' joins (equal keys must meet on one worker). Both of those are keyless or key-carrying in their own right, and the existing "Local exchange with parallelism requires keys" guard continues to protect the keyless one untouched. A UNION ALL input is simply never a keyless local exchange, so no new marker or discriminator is needed anywhere.

Degradation, and the co-residency fix

When the branches do not resolve to the same workers, the union stage cannot inherit both. Branches that still line up stay SINGLETON and wire 1-to-1; the rest are promoted to a hash shuffle on the projected columns. Correct either way for a concatenation, and the aligned branches still pay nothing.

The parallel wiring addresses a whole contiguous receiver range at the range's first host. That only holds when the receiver map was derived from the sender by WorkerManager#assignWorkersForLocalExchange. When the stage instead takes its workers from the candidate servers, a range spans several hosts — so blocks would be posted to a mailbox on the wrong server and the real receiver would block until the query deadline. This now verifies co-residency instead of assuming it, and falls back to a full hash shuffle when it does not hold. That is a fix for the pre-existing callers too, not only for UNION ALL.

A related consequence: a local exchange whose worker counts do not divide evenly used to throw. It now falls back to a hash shuffle, which is correct for every local-exchange kind since HashExchange routes each key consistently. A join hinted local on both sides with mismatched worker counts previously failed outright and now plans and runs correctly.

Set-op output distribution

PinotRelDistributionTraitRule had no SetOp case, so every set operation fell through to RANDOM_DISTRIBUTED. It now derives the output from what the input exchanges actually do:

• all inputs are hash exchanges (INTERSECT / EXCEPT / distinct UNION, or UNION ALL with the hint set to 'false') → hash distributed on all output columns, so a downstream exchange keyed on those columns can skip its own shuffle. This holds for is_colocated_by_set_op_keys='true' too: the hint asserts that rows equal across all projected columns already share a worker, which is exactly the claimed distribution, and we take it at its word just as the exchange does.
• any input is a local exchange (the UNION ALL case) → nothing claimed, because a local exchange redistributes nothing and therefore guarantees nothing. This is what keeps a deduplicating consumer — for example the aggregate UnionToDistinctRule puts over a distinct UNION — from skipping a shuffle it needs.

Latent bug fixes

All three were reachable before this change, two of them via an explicit is_colocated_by_set_op_keys='true' hint over fully-pruned inputs; the new default makes them reachable without a hint:

• MailboxAssignmentVisitor divided by zero when a leaf stage had all of its segments pruned.
• WorkerManager let a zero-worker child anchor the local-exchange assignment. A UNION ALL with one branch fully pruned would leave the union stage with no workers and silently return empty results.
• The parallel wiring mis-addressed receiver mailboxes whenever a contiguous range spanned hosts (see above).

Notes

• V1 only. The V2 physical optimizer (usePhysicalOptimizer=true) already plans UNION ALL without a shuffle and is untouched.
• No wire-format change: SINGLETON and HASH_DISTRIBUTED are pre-existing DistributionType values already handled by BlockExchange, and the decision is broker-side at plan time. No rolling-upgrade concern and no backward-incompat label.
• is_colocated_by_set_op_keys='true' is now a no-op on a UNION ALL (its inputs get a local exchange either way); only 'false' changes anything there. The hint is unreleased, so no compatibility shim is owed.
• There is deliberately no cluster-wide kill switch. Correctness never depends on this (concatenation is mapping-agnostic), the guards above prevent mis-wiring structurally, the per-query hint is the escape hatch, and V2 ships the same behaviour with no switch. Happy to add a broker config if reviewers would rather have one.
• Behaviour change: the default plan shape for every hint-free UNION ALL changes. Because the union stage inherits its inputs' worker layout instead of redistributing, input skew is carried into the union stage rather than rebalanced; in practice the next exchange above the union re-partitions.

Testing

• QueryCompilationTest — local exchange by default (with keys), the hint opt-out, misaligned branches where only the misaligned one is promoted to a shuffle, distinct set ops unaffected, and both directions of the distribution derivation.
• MailboxAssignmentVisitorTest — 1-to-1 wiring, promotion to a hash shuffle on unequal counts, per-host addressing when a receiver range spans servers, zero-sender wiring, and a negative control proving a keyless local exchange (the semi-join build side) still fails loudly.
• ExplainPhysicalPlans.json / SetOpPlans.json / AggregatePlans.json — updated plan shapes; the 'false' opt-out plan is retained.
• QueryHints.json (compared against H2, replayed on both optimizers) — UNION ALL correctness, dedup and GROUP BY above a local union, the colocated INTERSECT claim, and the mismatched-partition-count case.

Known gaps

Three of the four disagreement checks in canInheritWorkerAssignment (partition parallelism, partition classes, partition function) and the preserved "Found multiple local exchanges" rejection have no direct test. Neither is a known bug; both are places where a future change could regress quietly. Happy to add them here if preferred.

[Jackie-Jiang]

Consider not adding exchange from the beginning for UNION ALL. It is not really pre-partitioned

[codecov-commenter]

Codecov Report

❌ Patch coverage is 20.25316% with 63 lines in your changes missing coverage. Please review.
✅ Project coverage is 38.96%. Comparing base (f8352be) to head (a02289a).
⚠️ Report is 3 commits behind head on master.

Files with missing lines	Patch %	Lines
.../org/apache/pinot/query/routing/WorkerManager.java	16.66%	17 Missing and 3 partials ⚠️
...ery/planner/physical/MailboxAssignmentVisitor.java	5.88%	14 Missing and 2 partials ⚠️
...lcite/rel/rules/PinotRelDistributionTraitRule.java	0.00%	7 Missing and 1 partial ⚠️
...inot/calcite/rel/logical/PinotLogicalExchange.java	50.00%	6 Missing ⚠️
...te/rel/rules/PinotSetOpExchangeNodeInsertRule.java	0.00%	4 Missing ⚠️
.../pinot/query/planner/plannode/MailboxSendNode.java	33.33%	4 Missing ⚠️
...che/pinot/query/planner/plannode/ExchangeNode.java	40.00%	3 Missing ⚠️
.../query/planner/logical/EquivalentStagesFinder.java	0.00%	1 Missing ⚠️
.../query/planner/logical/RelToPlanNodeConverter.java	0.00%	0 Missing and 1 partial ⚠️

❗ There is a different number of reports uploaded between BASE (f8352be) and HEAD (a02289a). Click for more details.

HEAD has 4 uploads less than BASE

Flag	BASE (f8352be)	HEAD (a02289a)
java-25	6	5
temurin	6	5
unittests1	1	0
unittests	2	1

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #19325       +/-   ##
=============================================
- Coverage     67.00%   38.96%   -28.04%     
+ Complexity     1424     1423        -1     
=============================================
  Files          3463     3463               
  Lines        221671   221734       +63     
  Branches      34954    34975       +21     
=============================================
- Hits         148524    86402    -62122     
- Misses        61310   127454    +66144     
+ Partials      11837     7878     -3959     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	38.96% <20.25%> (-28.04%)	⬇️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	38.96% <20.25%> (-28.04%)	⬇️
unittests	38.96% <20.25%> (-28.04%)	⬇️
unittests1	?
unittests2	38.96% <20.25%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[yashmayya]

Superseded by #19330