yashmayya opened a new pull request, #18924: URL: https://github.com/apache/pinot/pull/18924
This addresses the remaining paths of https://github.com/apache/pinot/issues/18201. ## Overview #18237 added broker-side segment pruning to the MSE logical planner's **non-partitioned** leaf path. This PR extends it to the two remaining paths: * **Partitioned leaf path** (`assignWorkersToPartitionedLeafFragment`): the leaf filter is routed through the existing `getRoutingTable(PinotQuery)` mechanism (the same one the non-partitioned path uses), and a partition is dropped only when **every** one of its segments was pruned by the segment pruners. Surviving partitions keep their exact server assignment (`pickEnabledServer` is now indexed by `requestId + partitionId` instead of a running counter — byte-identical when nothing is pruned, and colocation-stable when partitions are skipped). * **Logical tables** (`assignWorkersToNonPartitionedLeafFragmentForLogicalTable`): the leaf filter is forwarded into each physical table's routing `BrokerRequest` (instead of the previous bare `SELECT *`), so the per-physical-table segment pruners can run. The pruned count is reported via `LogicalTableRouteInfo.getNumPrunedSegmentsTotal()`. Both paths report `numSegmentsPrunedByBroker` through the existing `DispatchablePlanContext` → `DispatchableSubPlan` → `BrokerResponseNativeV2` plumbing. Everything remains behind the same gate as #18237: the `useBrokerPruning` query option, defaulting to `pinot.broker.multistage.logical.planner.use.broker.pruning` (default **false**). With the gate off, behavior is unchanged (verified: the full `pinot-query-planner` suite of 1329 tests passes, and the empty-partition preconditions throw exactly as before when pruning is inactive). The same flag now gates broker pruning across **all** logical-planner leaf paths (non-partitioned from #18237, plus partitioned and logical tables here) — there is intentionally no per-path flag. Since it defaults off, ships in the same release as #18237, and pruning is a correctness-preserving routing optimization, enabling it simply extends the fan-out reduction to more query shapes. ### Design choice: route the filter, don't recompute partition ids An alternative considered was recomputing surviving partition ids directly from the filter literal via `PartitionFunctionFactory.getPartitionFunction(functionName, numPartitions, null)` (porting `SinglePartitionColumnSegmentPruner.isPartitionMatch`). This was rejected because `TablePartitionReplicatedServersInfo` only carries the partition function *name*, while functions like `Murmur3` consume `functionConfig` (seed, variant, normalizer) inside `getPartition()` — reconstructing with `null` config would silently compute wrong partition ids and drop live data on any table with a non-default partition function config. Routing the filter through the real segment pruners uses each segment's own partition metadata, which is correct for every function/config combination. ### Safety behaviors * **Fail-open everywhere**: unsupported leaf shapes, routing failures, or missing metadata fall back to unpruned routing. Broker pruning is a routing optimization only — servers still apply the filter — so the failure mode to avoid is dropping a segment that had matching rows, never scanning too much. * **All-pruned fallback**: if every partition is pruned, the leaf falls back to the unpruned assignment instead of producing an empty worker map. The all-leaves-empty short-circuit (#18538) only fires when *all* leaves are empty; in a multi-leaf plan (e.g. a join), a zero-worker leaf would leave the receiving exchange with no senders. The unpruned fallback keeps the plan well-formed and the server-side filter yields the correct empty result. * **Empty partitions**: with pruning active, a partition/worker with no segments is skipped instead of hitting the existing "Failed to find any segment" preconditions; with pruning off, the original throwing behavior is preserved exactly. ## Why colocated joins are NOT included (yet) The partitioned leaf path is also the colocated-join / pre-partitioned-exchange path, and pruning there can silently produce **wrong results** rather than just worse plans. The mechanism: * `MailboxAssignmentVisitor` wires a pre-partitioned sender to its receiver **1-to-1 by worker id** (`computeDirectExchange`). Worker id is the *only* channel carrying "which partition is this" between stages — nothing at runtime verifies that sender worker `i` and receiver worker `i` actually hold the same partition. * `isPrePartitionAssignment` / `isDirectExchangeCompatible` only compare worker **counts** and partition function names, never partition identity. * Pruning compacts worker ids (survivors are renumbered `0..k`). Two sides of a colocated join are routed independently (different tables, different filters, different segment/replica states), so they can prune to **different partition sets with coincidentally equal counts** — e.g. side A keeps partitions {1,3}, side B keeps {0,2}. The count checks pass, the direct exchange pairs A's partition 1 with B's partition 0, and the join returns silently wrong results. So pruning is gated on the exchange type: it engages only where the leaf's output is **shuffled** (`connectWorkers` re-hashes across any worker count — provably safe under compaction), plus one carve-out described next. ### The safe pre-partitioned carve-out: single derived chain to the root A pre-partitioned leaf **is** pruned when its output flows to the root exclusively through single-input, scan-free intermediate stages (`isSingleDerivedChainToRoot`). This covers the common `GROUP BY <partition key>` / window-over-partition-key shapes. It is safe because: 1. Each stage on such a chain derives its worker map *from* the (compacted) leaf via `assignWorkersForLocalExchange`, so the 1-to-1 pairing follows the compaction automatically, and per-key operators stay correct (equal keys never split across workers). 2. The chain never meets another plan branch below the root, so the compacted worker count cannot coincidentally re-enable a direct exchange against an unrelated pre-partitioned sibling. (This is why the condition walks all the way to the root rather than just checking the immediate parent: a compacted count can propagate through derived stages and collide with a sibling branch several exchanges up.) Stages containing a table scan (e.g. the probe side of a colocated dynamic-broadcast semi-join, which is a leaf *containing* a join) and multi-receiver sends (spooled stages) are excluded — their worker maps are not derived from the pruned chain. ## How the v2 physical optimizer handles this, and why v1 doesn't just copy it The v2 optimizer (`LeafStageWorkerAssignmentRule` + `WorkerExchangeAssignmentRule`) already supports broker pruning including partitioned tables. Its approach differs structurally: * **Workers are per-server, not per-partition**: partition `P` maps to worker `P % numWorkers` with `numPartitions` taken from **table metadata**, never from the surviving set. A pruned partition just contributes zero segments to its slot. Worker identity is therefore stable across both sides of a join regardless of uneven pruning — colocation and pruning coexist in the common case. * **Exchanges are derived from computed facts, not hints**: worker assignment happens *during* physical planning, and the decision to elide a shuffle compares actual `PinotDataDistribution`s — worker count, worker hash (the concrete `workerId@instance` list), hash function, and partition count. If pruning empties an entire server on one side, the mod-mapping invariant breaks, the side degrades to unpartitioned distribution, and a shuffle is inserted. Worst case is a slower plan, never a wrong one. **Tradeoffs of porting that approach to the v1 logical planner:** * v1's pipeline is one-way: plan shape and `isPrePartitioned` flags are frozen first, `WorkerManager` assigns workers second, `MailboxAssignmentVisitor` wires mailboxes third with only count-parity checks. There is no channel to feed "this side pruned partitions {1,3}" back into the exchange decision. Retrofitting one is a redesign of the v1 planning pipeline — at which point the community-preferred path is v2 feature parity (#15455) rather than backporting v2's architecture. * Keeping v2-style empty worker slots in v1 would preserve alignment but defeat the purpose: an empty worker is still a dispatched worker (gRPC, thread, mailboxes on that server), so `numServersQueried` — the motivating metric in #18201, where one gray-failing server degrades every query — doesn't improve. The fan-out win requires *not dispatching*, which requires compaction, which requires the alignment safety gate above. * Sparse worker ids (drop the worker, keep the numbering) are not representable: mailbox wiring and `WorkerMetadata` assume dense `0..N-1` ids. The chosen design (compaction where shuffled or single-derived-chain, abstention on multi-branch pre-partitioned shapes) captures the fan-out reduction for non-colocated queries and partition-key group-bys, while colocated joins conservatively stay correct-but-unpruned. Extending pruning to colocated joins would need either partition-aware exchange wiring in v1 or join-key partition inference (prune the same partitions on both sides provably) — both are follow-up-sized efforts. ## Testing Unit (`WorkerManagerTest`, partitioned path): * equality prune (count computed from dropped partitions, cross-checked against a deliberately-wrong routing-reported count), multi-partition keep, all-pruned → unpruned fallback, disabled no-op, multiple-partitions-per-worker; * `GROUP BY partition key` prunes via the derived-chain carve-out, and the same shape with `ORDER BY ... LIMIT` prunes through a multi-hop derived chain; * colocated self-join stays gated (all partitions assigned, count 0) — proving the gate discriminates, since the identical filter prunes in the group-by case; * a partition kept alive by an *unavailable* (not pruned) segment, so transient outages never drop data; * routing-failure fail-open (getRoutingTable throws → unpruned assignment, query still planned); * server-placement stability — with multiple replicas per partition, pruning partition 0 does not shift the surviving partitions' server assignments (locking in the `requestId + partitionId` seeding; the test was confirmed to fail if reverted to a running counter). Unit (logical table): `buildLogicalTableRoutingBrokerRequest` filter/typed-name/options forwarding, `SELECT *` fallback, and non-mutation of the source query. Integration: * `SegmentPartitionLLCRealtimeClusterIntegrationTest.testMultiStageBrokerPruningOnPartitionedTable` — a **real** partition pruner drives partition survival end to end: an MSE partition-key query with `useBrokerPruning=true` reports `numSegmentsPrunedByBroker > 0`, queries strictly fewer segments than the unpruned run, and returns identical results. * `BaseLogicalTableIntegrationTest` — result parity with pruning on vs. off across both query engines (smoke coverage of the logical `calculateRoutes` path with a forwarded filter). -- 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]
