yashmayya opened a new pull request, #19330:
URL: https://github.com/apache/pinot/pull/19330
`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
They 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.`
It also means no new marker or discriminator is needed anywhere. `SINGLETON`
is shared with the colocated dynamic-broadcast semi-join build side, which is
*keyless* and must never be redistributed (every receiver needs the whole build
side — the non-colocated variant broadcasts for exactly that reason). Because a
`UNION ALL` input always carries keys, the existing `"Local exchange with
parallelism requires keys"` guard keeps protecting that path untouched.
### Worker assignment
`WorkerManager` previously rejected a stage with more than one
local-exchange child outright, which every multi-branch `UNION ALL` would hit.
It now accepts them when they all agree on worker map, partition parallelism,
partition function and partition classes — the stage inherits from the first,
so it may only do so when that one describes them all. A child with no worker
at all (a fully pruned leaf) can never anchor it, or the stage would end up
with no workers and silently drop every live sibling's rows.
When they disagree, the stage falls back to the regular assignment and each
branch is promoted to a hash shuffle. Branches that still line up 1-to-1 stay
local, so a misaligned sibling does not cost the aligned ones anything.
### Co-residency fix
The parallel wiring in `computeDirectExchangeWithParallelism` 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
`assignWorkersForLocalExchange`. When the stage instead takes its workers from
the candidate servers, a range spans several hosts — 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
hash shuffle when it does not hold. That is a fix for the existing callers too,
not only for `UNION ALL`.
A 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.
Note the 1-to-1 path is unaffected: it already resolves the server per
worker and emits separate mailboxes when sender and receiver diverge, so
aligned-but-not-co-located stays correct and merely loses the in-process
handover.
### 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.
This part is independently useful and could be split into its own PR if
reviewers prefer; it is not required for the `UNION ALL` change, since set ops
already defaulted to `RANDOM_DISTRIBUTED`.
### 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, 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.
- One small unrelated fix rides along because this feature's own escape
hatch can reach it: `isDirectExchangeCompatible` divided by zero when a leaf
stage had all of its segments pruned.
### Testing
- `QueryCompilationTest` — local exchange with keys by default, 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` — 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) have no direct
test. Not a known bug, but a place a future change could regress quietly. Happy
to add them here if preferred.
--
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]
