yashmayya opened a new pull request, #19325:
URL: https://github.com/apache/pinot/pull/19325
`UNION ALL` is a pure concatenation, so any row-to-worker mapping produces
the same result. 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"), so it never constrains the
distribution of a union's inputs. 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 with V2. When no `is_colocated_by_set_op_keys` hint
is present, `UNION ALL` input exchanges now default to pre-partitioned (direct,
1-to-1) sends.
### What it buys
Once the exchange is marked pre-partitioned, the existing machinery does the
rest:
- `WorkerManager#isPrePartitionAssignment` lets the union stage inherit its
leaf stages' workers, which keeps the stages above it aligned as well — the
main win for plans with many colocated stages.
- `MailboxAssignmentVisitor#computeDirectExchange` wires sender worker *i*
to receiver worker *i*. When both land on the same server they share an
in-memory mailbox, so there is no network hop and no block serialization.
- `HashExchange#route` short-circuits when there is a single destination
mailbox, so even the per-row hashing disappears (the plan still reads
`HASH_DISTRIBUTED`).
### Why it cannot mis-wire a plan
Correctness does not depend on placement here, and a direct exchange is only
formed when the shapes line up: `isPrePartitionAssignment` and
`isDirectExchangeCompatible` both require equal worker counts and matching
partition functions across branches. When they do not match (say inputs
partitioned 2 ways and 4 ways), the planner falls back to today's shuffle, so
the worst case is a missed optimization rather than a bad plan. There is a
runtime test for that fallback.
Opt out per query with `/*+
setOpOptions(is_colocated_by_set_op_keys='false') */`.
### Set operation output distribution
`PinotRelDistributionTraitRule` had no `SetOp` case, so every set operation
fell through to `RANDOM_DISTRIBUTED`. It now derives the output distribution
from what the input exchanges actually do:
- every input genuinely shuffles (`INTERSECT` / `EXCEPT` / distinct `UNION`,
or `UNION ALL` with the hint set to `'false'`) → hash distributed on all output
columns, which lets a downstream exchange keyed on those columns skip its own
shuffle;
- any input pre-partitioned → no distribution claimed, because the mailbox
layer may still fall back to a real shuffle.
### Two latent bug fixes
Both were already reachable before this change via an explicit
`is_colocated_by_set_op_keys='true'` hint over fully-pruned inputs. The new
default makes them reachable without a hint, so they are fixed here:
- `MailboxAssignmentVisitor#isDirectExchangeCompatible` divided by zero when
a leaf stage had all of its segments pruned (0 senders and 0 receivers pass the
multiplication check).
- `WorkerManager#isPrePartitionAssignment` let a zero-worker child anchor
the local-exchange assignment, because `partitionCount == 0` conflated "unset"
with "computed zero". A `UNION ALL` with one branch fully pruned would then
leave the union stage with no workers and silently return empty results.
### Notes
- V1 only. The V2 physical optimizer (`usePhysicalOptimizer=true`) already
plans `UNION ALL` without a shuffle and is untouched.
- This changes the default plan shape for hint-free `UNION ALL` queries.
Because the union stage now 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.
- Not wire-incompatible: pre-partitioned exchanges are decided broker-side
at plan time, so there is no rolling-upgrade concern and no `backward-incompat`
label.
### Testing
- `QueryCompilationTest` — pre-partitioned by default, the hint opt-out,
distinct set operations unaffected, and both directions of the new distribution
derivation.
- `ExplainPhysicalPlans.json` — the new default plan shape, with the
previous shuffled plan retained under the hint.
- `QueryHints.json` (compared against H2, replayed on both optimizers) —
`UNION ALL` correctness, dedup and `GROUP BY` above a pre-partitioned union,
and the mismatched-partition-count fallback.
- `MailboxAssignmentVisitorTest` / `WorkerManagerTest` — regression tests
for the two fixes above.
--
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]
