Support colocated joins when a partition holds no segments

[yashmayya]

Why

A colocated join fails today when a partition of one of its tables holds no segments:

Failed to find any segment for table: myTable_OFFLINE, partition: 3

This is easy to hit and hard to avoid:

• A partition key with skew leaves some partitions with no rows.
• Retention removes every segment of a partition.
• A stream partition exists but has never committed a segment.
• One logical partition space is split across several tables, so each table populates a subset.
• A table declares more partitions than it currently uses, to leave room to grow.

The assignment cannot simply skip the empty partition. Worker ids come from a counter over the partitions that hold data, so skipping one shifts every later partition down a slot. Two tables that each drop a different empty partition then end up with equal worker counts, get wired one-to-one, and join mismatched partitions. That loses rows and reports no error, which is why the assignment refuses to continue at all.

What changes

• The stages tied together by direct exchanges form a group, and the group shares one ordered list of partition classes.
• A class leaves the list only when no member of the group holds data in it.
• A class the group keeps, but one member holds no data for, gets a worker with no segments. That worker sits on a server borrowed from a member that does hold the class, so the exchange stays in process.
• Both sides of every direct exchange assert that they agree on the list. A disagreement fails at plan time instead of returning wrong rows.
• The broker now publishes the partitions whose only segments are new and have no online replica. Those hold data that no server can serve as a whole, so they keep failing.

A side effect is less fan-out. A table that declares 8 partitions and populates 3 runs 3 workers, so the broker sends fewer requests and waits on fewer servers.

What now fails on purpose

• A hybrid table with segments whose partition metadata does not match the table config now fails planning. It used to answer and drop those segments in silence. The most common cause is a partition config added after some segments were built.
• A table spread over more than one cluster reports no partition info, so a partition-aware plan is not attempted on a partial view. No server holds another cluster's segments, so one cluster's view would make a partition served elsewhere look empty.

Broker pruning comes next

Broker pruning is off for a colocated join today, and this change is what makes it possible.

• Pruning drops partitions that hold no matching segment. That is the same operation as dropping a class that holds no data.
• The group already shares one class list, so a pruned class can leave the list the same way an empty one does, and both sides stay in step.
• The gain is fewer servers per query, which protects tail latency on large clusters.

A follow-up PR adds it.

Testing

• 19 files. Planner, runtime and broker suites all pass.
• An end-to-end test joins two tables that declare 8 partitions and populate 3 and 4 of them. It asserts the rows, the worker count per leaf, the segments each leaf reads, and that each send goes to one receiver with the same worker id. That last check proves the plan did not fall back to a shuffle.
• One test drives the broker metadata directly and proves the new signal does not depend on map iteration order.

Two things left out on purpose

Aggregation merge identity is not covered here. A leaf that scans nothing emits one identity row for an aggregation with no GROUP BY. That predates this change: a worker whose segments are all pruned on the server already does the same. This change raises how many such rows reach the merge without introducing the dependency. The one aggregation that is not a true merge identity, PERCENTILEKLL, fails only when every worker is empty, which this change does not newly reach. It is tracked separately.

A worker with no segments is charged against the query thread estimate like any other worker. It is dispatched and does run a leaf operator. The estimate over-counts by two threads per such worker. That is conservative, and it only affects colocated joins over a partition space that is largely unpopulated. Those queries failed outright before, so there is no earlier estimate to compare against.

Labels

Applied: bug, backward-incompat, release-notes, multi-stage, query

[codecov-commenter]

Codecov Report

❌ Patch coverage is 91.86047% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.68%. Comparing base (14333f3) to head (45447e7).
⚠️ Report is 15 commits behind head on master.

Files with missing lines	Patch %	Lines
.../org/apache/pinot/query/routing/WorkerManager.java	90.21%	4 Missing and 19 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 ⚠️
...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   #19166      +/-   ##
============================================
+ Coverage     66.57%   66.68%   +0.11%     
  Complexity     1423     1423              
============================================
  Files          3441     3445       +4     
  Lines        218320   219011     +691     
  Branches      34737    34886     +149     
============================================
+ Hits         145341   146046     +705     
+ Misses        61251    61226      -25     
- Partials      11728    11739      +11     

Flag	Coverage Δ
custom-integration1	?
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	66.68% <91.86%> (+0.11%)	⬆️
temurin	66.68% <91.86%> (+0.11%)	⬆️
unittests	66.68% <91.86%> (+0.11%)	⬆️
unittests1	57.24% <91.54%> (+0.13%)	⬆️
unittests2	38.89% <25.58%> (+<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.

[Copilot]

Pull request overview

Enables colocated joins to handle empty partitions while preserving partition-to-worker alignment and reducing fan-out.

Changes:

• Adds shared partition-class assignment and empty-worker padding.
• Adds deferred-partition and multi-cluster safeguards.
• Expands planner, broker, mailbox, and integration coverage.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
pinot-query-planner/.../WorkerManagerTest.java	Covers partition assignment and padding.
pinot-query-planner/.../ColocationGroupAnalyzerTest.java	Tests reducible group detection.
pinot-query-planner/.../QueryEnvironmentTestBase.java	Updates partition metadata construction.
pinot-query-planner/.../PinotDispatchPlannerTest.java	Tests spooled-stage assignment.
pinot-query-planner/.../MailboxAssignmentVisitorTest.java	Tests direct-exchange class agreement.
pinot-query-planner/.../DispatchableSubPlanTest.java	Tests segment-map preservation.
pinot-query-planner/.../WorkerManager.java	Implements class reduction and padding.
pinot-query-planner/.../LeafPartitionHints.java	Centralizes partition-hint parsing.
pinot-query-planner/.../ColocationGroupAnalyzer.java	Identifies colocated fragment groups.
pinot-query-planner/.../MailboxAssignmentVisitor.java	Validates direct-exchange mappings.
pinot-query-planner/.../DispatchablePlanMetadata.java	Stores broker-local class metadata.
pinot-query-planner/.../DispatchablePlanFragment.java	Preserves segment assignments when copied.
pinot-query-planner/.../DispatchablePlanContext.java	Caches partition metadata per query.
pinot-integration-tests/.../ColocatedJoinEmptyPartitionTest.java	Adds end-to-end join coverage.
pinot-core/.../TablePartitionReplicatedServersInfo.java	Publishes deferred partitions.
pinot-broker/.../SegmentPartitionMetadataManagerTest.java	Tests deferred-partition detection.
pinot-broker/.../MultiClusterRoutingManagerTest.java	Tests cross-cluster partition metadata.
pinot-broker/.../SegmentPartitionMetadataManager.java	Tracks deferred-only partitions.
pinot-broker/.../MultiClusterRoutingManager.java	Avoids combining partition arrays.

Suppressed comments (2)

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java:239

• A failed remote lookup leaves it unknown whether that cluster also serves the table. Continuing can return another cluster's partition array and enable a partition-aware plan over a partial view, which can silently drop remote-only partitions. Return null on this failure so the caller falls back instead.

      } catch (Exception e) {
        LOGGER.error("Error getting table partition info from remote cluster routing manager for table {}",
            tableNameWithType, e);
        continue;

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java:243

• A null partition-info response does not mean this cluster lacks the table; a routing entry can exist without usable partition metadata. Treating that cluster as absent can return the local array for a table spread across clusters and silently omit its partitions. Check routingExists and return null when the remote cluster serves the table.

      if (remotePartitionInfo == null) {
        continue;
      }