fix: handle optional segments, unavailable segments, and empty partitions in WorkerManager (#18223)

[waterWang]

Closes #18223

[Copilot]

Pull request overview

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Updates WorkerManager routing/assignment logic to better tolerate optional segments, unavailable segments, and partitions with no segments, addressing #18223.

Changes:

• Merge required + optional segments when building server→segment mappings for multi-stage execution.
• Propagate unavailable segments into dispatchable metadata for replicated leaf stages.
• Skip (instead of failing) partitions that have no segments, treating them as pruned/empty.

[yashmayya]

Thanks for picking this up. Unfortunately I don't think this can be merged in its current form — I traced each of the four changes through to its runtime behaviour, and the common problem is that these four TODOs are not unimplemented features but load-bearing guards. Removing them without building the mechanism each one is waiting for turns loud, correct failures (and one documented, bounded gap) into either a new hard failure on a routine steady-state condition, or silently wrong results.

Line references below are against master at 863b9da.

1. Merging optional segments into the required list inverts the meaning of the field

This affects both assignWorkersToNonPartitionedLeafFragment and transferToServerInstanceLogicalSegmentsMap.

A segment is classified as optional precisely when the selected replica is not ONLINE. From BalancedInstanceSelector.java:77:

// This can only be offline when it is a new segment. And such segment is marked as optional segment so that broker or server can skip it upon any issue to process it.

"New" here means younger than pinot.broker.new.segment.expiration.seconds, default 5 minutes (CommonConstants.java:590).

The two-list split is the feature. BaseTableDataManager.acquireSegments (BaseTableDataManager.java:744-765) walks the two lists differently, and the comment at line 758 states the contract:

// Optional segments are not counted to missing segments that are reported back in query exception.

Required segments that fail to acquire land in missingSegments; optional ones are skipped silently. So the question is what the merged list is treated as on the server side, and on both MSE leaf paths the answer is required:

• ServerPlanRequestUtils.compileInstanceRequest (ServerPlanRequestUtils.java:213) never calls instanceRequest.setOptionalSegments(...).
• The logical-table variant only calls tableSegmentsInfo.setSegments(...) (ServerPlanRequestUtils.java:430).

There is no optional-segment channel in the MSE stage metadata at all, so merging the lists at the routing site does not add support for optional segments — it promotes them to required.

The consequence chain:

1. Required-but-absent segment → ServerQueryExecutorV1Impl.java:298 adds a SERVER_SEGMENT_MISSING exception.
2. In MSE that exception is fatal, per the comment at LeafOperator.java:544: "Currently MERGE_RESPONSE_ERROR and SERVER_SEGMENT_MISSING_ERROR are counted as query failure." (SSE tolerates it — BrokerReduceService.java:89 can even filter it out of the response — but MSE fails the whole query.)

So the net effect of this hunk: any MSE query touching a table that has a segment created in the last 5 minutes which is not yet ONLINE on the selected replica now fails outright. That is a routine condition — it occurs after every realtime segment commit and after every offline push. Before this change the same query succeeded, missing at most a few of the very newest rows. That bounded freshness gap is what the TODO is acknowledging; the diff replaces it with a hard failure.

It is also racy rather than deterministic: whether the server happens to have already loaded the segment decides whether the query fails, so in production this surfaces as intermittent, timing- and load-dependent query failures rather than a reproducible bug.

Two smaller defects in the same hunks:

• optionalSegments != null is dead code. The only production construction site is BaseBrokerRoutingManager.java:1125, which always passes a non-null ArrayList.
• Wrapping only the merged branch in Collections.unmodifiableList makes the list's mutability data-dependent — the same list is the raw mutable routing list when no optional segments exist. filterLeafStageSegments / filterReplicatedLeafStageSegments are explicitly filtering extension points, so this is a latent UnsupportedOperationException that only fires when optional segments happen to be present.

2. Skipping empty partitions breaks the partition-to-workerId identity, silently

This one is documented directly above the changed code. computePartitionsToKeep (WorkerManager.java:983-1015) already implements worker dropping for broker pruning, and its javadoc enumerates the two cases where it deliberately refuses to:

- the leaf feeds a pre-partitioned (1-to-1 direct) exchange -- dropping/compacting workers would misalign sender/receiver worker ids in MailboxAssignmentVisitor.
- every partition would be pruned -- an empty worker map would break exchanges in a multi-leaf plan

enforced by the metadata.isPrePartitioned() bail-out at WorkerManager.java:1013. The same invariant is stated again in MailboxAssignmentVisitor.computeDirectExchange (MailboxAssignmentVisitor.java:105), which explains why the cross-server fallback is safe:

// back to a cross-server send: the exchange stays correct because worker id still maps to the same partition on both sides

The new continue in assignOnePartitionPerWorker compacts workers with no isPrePartitioned gate, so it breaks exactly that invariant. Note that workerId is a compacted running counter (workerId++) while the server is picked from requestId + i; skipping a partition therefore preserves the server assignment but shifts the worker id to partition mapping.

Concrete wrong-results case — colocated join A ⋈ B, 4 partitions each, one worker per partition. A has segments in partitions {0, 2, 3}; B has segments in {0, 1, 3}. Both sides compact to 3 workers with the same partition function and the same partition count, so isPrePartitionAssignment (WorkerManager.java:339-372) returns true and a direct 1-to-1 exchange is chosen. A's worker 1 holds partition 2 and B's worker 1 holds partition 1, and they are joined against each other. No exception and no user-visible warning — only the broker-side LOGGER.warn added by this diff, which nothing correlates with the wrong result. Because the mapping now depends on which partitions happen to be empty at plan time, the same query silently returns different wrong answers as data lands and as retention drops segments.

Separately, the premise that partitionInfo == null means "this partition contains no segments" is not correct. See SegmentPartitionMetadataManager.computeTablePartitionReplicatedServersInfo (SegmentPartitionMetadataManager.java:276-287): a partition whose only segments are new (<5 min) and do not yet have their replicas available is left null via excludedNewSegments. The new continue therefore silently drops real, imminently-queryable data from the result instead of failing, which is a correctness regression, not a resiliency improvement.

It also does not fix the failure this is most often reported as. A partition whose segments are genuinely unavailable produces a non-null PartitionInfo with an empty _fullyReplicatedServers (SegmentPartitionMetadataManager.java:223-233), so it still trips the "Failed to find enabled fully replicated server" precondition two lines below the one that was removed.

Finally, the change is internally inconsistent with the sibling path: assignMultiplePartitionsPerWorker (WorkerManager.java:1152-1159) keeps the equivalent guard. After this diff, the identical cluster state throws when partitionSize yields more than one partition per worker and silently skips when it yields exactly one.

On the replacement comment — "The leaf stage can handle empty partitions by returning an empty response, which is equivalent to a pruned partition" is asserted without demonstration, and it is also beside the point: the diff does not create an empty worker, it removes the worker. The blocker was never the empty response, it is worker-id alignment.

3. Attaching unavailable segments in setSegmentsForReplicatedLeafFragment

This hunk computes routing a second time, and the result it reads is not the routing the dispatched segment list was built from. The replicated path's segments come from getSegments(...) → RoutingEntry.getSegments (BaseBrokerRoutingManager.java:1461-1471), which applies the segment selector and the segment pruners and never consults the instance selector. Consequences:

• False positives. Segments reported as unavailable are still dispatched — setReplicatedSegments(segmentsMap) is unchanged — and on a fully-replicated table will usually be served correctly. Each one nonetheless produces a user-visible SERVER_SEGMENT_MISSING exception on the response (MultiStageBrokerRequestHandler.java:816-829). This adds errors to healthy queries.
• Query options are dropped. The segment list is built with context.getPlannerContext().getOptions(), which honours the table sampler; the new call uses the getRoutingTable(tableName, requestId) overload, which passes Map.of() (WorkerManager.java:615-617). The two views therefore disagree about which segments are even in scope, so the loop can report segments that are not in the dispatched set at all.
• Hybrid mismatch. The loop is placed after the segmentsMap.remove(TableType.OFFLINE.name()) at WorkerManager.java:736, but iterates the freshly computed map, which still contains OFFLINE. It will report unavailable OFFLINE segments for a query that is not reading the offline side.
• Cost. It adds a full extra routing-table computation — per-segment instance selection over the whole table — to the planning path for every replicated leaf stage on every query, while the surrounding code already has the segment information it needs.

There is also an unresolved semantic question underneath this TODO: for a fully replicated table whose leaf stage runs on servers that host every segment, it is not clear that "no selected server hosts this segment" is a meaningful error condition at all. That needs to be settled before it is surfaced to users as a query exception.

4. Scope and coverage

The diff is 34 added lines in one file, but it is four independent behavioural changes, four deleted TODOs and one deleted Preconditions guard, with no tests. Two of the four are silent-wrong-result classes that only a test would catch. The linked issue (#18223) explicitly scopes "Add planner and runtime tests for replicated and partitioned leaf stages" and "Implement the missing routing and metadata handling, or fail with clearer behavior if a case must stay unsupported" — the current diff does neither for these cases, and would be much easier to review split one case per PR.