Support broker segment pruning for colocated joins

[yashmayya]

Follow-up to #19166, which reduced a colocated join to the partition classes that hold data. This does the same for the classes that a filter empties, so a filtered colocated join runs on fewer workers and reaches fewer servers.

Only the default (logical) MSE planner is affected. The V2 physical optimizer already derives its partitions from a pruned routing table.

What this enables

A colocated join over two partitioned tables, with a filter on the partition key:

SELECT /*+ joinOptions(is_colocated_by_join_keys='true') */ l.key, r.value
FROM  left  /*+ tableOptions(partition_key='key', partition_size='8', partition_function='Modulo') */ l
JOIN  right /*+ tableOptions(partition_key='key', partition_size='8', partition_function='Modulo') */ r
ON l.key = r.key
WHERE l.key IN (1, 2) AND r.key IN (1, 2)

Both sides prune every partition class except 1 and 2. The group drops the rest. The result:

• Fewer workers. Each leaf runs 2 workers instead of 8.
• Fewer servers. The broker dispatches only to the servers that hold classes 1 and 2. It also waits on only those servers, so the query loses the tail latency of the servers it skips.
• Fewer segments. Each worker scans only its own class.
• A cheaper cancel. Cancellation reaches the same reduced server set.

These cases get the reduction:

• A filter on the partition key, written on both sides.
• A filter on the partition key on one side, when Calcite carries it to the other side through the join equality.
• A partition class that neither table populates. Support colocated joins when a partition holds no segments #19166 already dropped these.

The 1-to-1 exchange survives the reduction. Every leaf of the group keeps one worker per surviving class, in the same order, so worker k means the same class on every side.

What this does not help

The common case is in this list, so it is worth stating plainly.

• A time filter. The time pruner removes segments, not partitions. A table partitioned by key and segmented by time holds every time range in every partition. This is the most common filter in Pinot, and it buys nothing here.
• A filter on one side only. A class survives while any member of the group still holds a matching row. An unfiltered side keeps every class it populates. See the trade-off below.
• A table without segmentPrunerTypes: ["partition"]. Pinot builds the partition pruner only when the routing config asks for it. Without it, nothing is pruned.
• A filter that matches nothing. The group keeps every populated class. A group with zero workers has nothing to wire a 1-to-1 exchange to, and the servers return the same empty result anyway.
• The probe leaf of a colocated semi-join. That leaf holds the join, so the routing-query builder cannot fold it.
• A non-colocated stage above the join, with the default config. getCandidateServersPerTables ignores the query and expands the dispatched set again. With useLeafServerForIntermediateStage=true, the reduction carries upward.

Trade-offs

A class drops only when every member of the group prunes it. The alternative is to drop a class as soon as one side prunes it. That is wrong for a RIGHT join, a FULL join, a union, and an anti-join. In each of those, the other side still produces rows. The union rule needs no knowledge of the operator above. A new colocated operator therefore cannot break it in silence. The cost is the one-sided filter case above.

A surviving class dispatches all of its segments on every member. A member whose own filter excludes them still scans them. The server-side filter removes them again. This keeps the class on the same server it uses without pruning, so the exchange stays in process. A finer verdict inside a surviving class is possible later.

Planning costs one routing call per member. Only a leaf that carries a filter pays it. The result is cached per fragment, so the leaf assignment reuses it.

A sound proof of emptiness

The old verdict read "this partition is pruned" from the absence of its segments in the routing table. Absence has innocent causes. Instance selection can class a segment as optional. The server that holds a segment can leave the enabled server map. A segment can enter the partition metadata before it becomes selectable. Each of these loses matching rows.

This adds a planner-only entry point that reports which segments the pruners pruned:

// RoutingManager, added as a default that proves nothing, so no implementation breaks
@Nullable Set<String> getPrunedSegments(BrokerRequest brokerRequest);

Only presence in that set is a proof. Instance selection takes no part and there is no request id, so two leaves over one table and one filter cannot disagree. MultiClusterRoutingManager combines the set by intersection, because a proof holds only when every cluster that can route the segment pruned it.

The ordinary partitioned leaf now uses the same proof. That path is strictly more conservative than before.

How to turn it off

No new flag. The query option useBrokerPruning and the broker config pinot.broker.multistage.logical.planner.use.broker.pruning already do it. With pruning off, every populated class keeps its worker, which is the behavior of #19166.

Tests

• WorkerManagerTest covers both sides pruning and one side pruning. It also covers the all-pruned fallback, an unavailable segment, a routing failure and the kill switch. Two more cases cover padding together with reduction, and a self-join whose sides prune different classes.
• BrokerRoutingManagerTest and MultiClusterRoutingManagerTest cover the pruned set and the multi-cluster intersection.
• ColocatedJoinEmptyPartitionTest runs the filtered join end to end. It proves the 1-to-1 wiring at the reduced width. It also matches the rows against the same query without the colocation hint. A filter that matches nothing returns an empty result.

The integration fixture had no RoutingConfig, so it built no partition pruner. This adds one. Without it a filtered assertion proves nothing.

Notes

• Depends on Support colocated joins when a partition holds no segments #19166.
• The dispatched server count does not move in the integration cluster. It runs 2 servers with numReplicas 2, so both servers hold every segment. The end-to-end tests assert worker count, segments queried, fan-out and receiver ids. WorkerManagerTest asserts the dispatched server set.
• A worker with no segments on an upsert table can still acquire newly-added segments, so it does not always scan nothing. This predates the change and is filed separately.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.41379% with 44 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.69%. Comparing base (14333f3) to head (ec4ac46).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
.../org/apache/pinot/query/routing/WorkerManager.java	91.07%	6 Missing and 23 partials ⚠️
...e/routing/TablePartitionReplicatedServersInfo.java	42.85%	3 Missing and 1 partial ⚠️
...e/pinot/query/routing/ColocationGroupAnalyzer.java	96.33%	0 Missing and 4 partials ⚠️
...oker/routing/manager/BaseBrokerRoutingManager.java	89.65%	1 Missing and 2 partials ⚠️
...apache/pinot/query/routing/LeafPartitionHints.java	93.10%	0 Missing and 2 partials ⚠️
...mentpartition/SegmentPartitionMetadataManager.java	91.66%	0 Missing and 1 partial ⚠️
...ery/planner/physical/MailboxAssignmentVisitor.java	90.90%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19174      +/-   ##
============================================
+ Coverage     66.57%   66.69%   +0.12%     
  Complexity     1423     1423              
============================================
  Files          3441     3445       +4     
  Lines        218320   219061     +741     
  Branches      34737    34910     +173     
============================================
+ Hits         145341   146100     +759     
+ Misses        61251    61218      -33     
- Partials      11728    11743      +15     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	66.69% <92.41%> (+0.12%)	⬆️
temurin	66.69% <92.41%> (+0.12%)	⬆️
unittests	66.69% <92.41%> (+0.12%)	⬆️
unittests1	57.22% <92.04%> (+0.11%)	⬆️
unittests2	38.93% <28.62%> (+0.03%)	⬆️

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.