Skip block copies for spool stages whose workers are all remote

[yashmayya]

PR flow

BroadcastExchange now decides block copying based on whether any destination mailbox delivers blocks by reference, avoiding unnecessary copies for spool stages with all-remote workers.

flowchart TD
  N0["BroadcastExchange#35;route #40;F6#41;"]:::stModified
  N1["BlockExchangeSendingMailbox#35;deliversByReference #40;F5#41;"]:::stModified
  N2["BlockExchange#35;#95;deliversByReference #40;F5#41;"]:::stAdded
  N3["SendingMailbox#35;deliversByReference #40;interface#41; #40;F3#41;"]:::stAdded
  N4["GrpcSendingMailbox#35;deliversByReference #40;F1#41;"]:::stAdded
  N5["InMemorySendingMailbox#35;deliversByReference #40;F2#41;"]:::stAdded
  N0 -->|"calls"| N1
  N1 -->|"returns"| N2
  N2 -->|"anyDeliversByReference calls"| N3
  N3 -->|"implemented by"| N4
  N3 -->|"implemented by"| N5
  classDef stAdded fill:#dafbe1,stroke:#1a7f37,color:#1f2328,stroke-width:2px
  classDef stModified fill:#fff8c5,stroke:#9a6700,color:#1f2328,stroke-width:2px
  classDef stRemoved fill:#ffebe9,stroke:#cf222e,color:#1f2328,stroke-width:2px
  classDef stUnchanged fill:#f6f8fa,stroke:#656d76,color:#1f2328,stroke-width:1px

Loading

AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing

Diff evidence

• F1: pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java — before · after
• F2: pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java — before · after
• F3: pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/SendingMailbox.java — before · after
• F5: pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java — before · after
• F6: pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BroadcastExchange.java — before · after

• Regenerate PR flow

Fixes #19427. Follow-up to #19353.

Problem

#19353 made BroadcastExchange copy blocks that carry aggregation intermediate results, because a destination can hand an on-heap block to a receiver by reference. It skips the copy for destinations that are not local, because those serialize the block within send.

That skip never applied to spools, which are the case #19353 was written for. A multi-send (spool) node wraps each receiver stage's inner exchange as a BlockExchange.BlockExchangeSendingMailbox, and isLocal() on that wrapper returns true unconditionally. Every outer destination of a spool is such a wrapper, so a receiver stage with no worker on the sending server still got a copy. The inner exchange then serialized that copy immediately, on the same thread, and no receiver ever mutated it. Each wasted copy costs one serialization plus one deserialization of every non-null OBJECT cell.

The cause is that isLocal() answered two different questions:

• BlockExchange#sendBlock asks "can I pass the block whole, without splitting it?"
• BroadcastExchange#route asks "can a receiver keep a reference to this block?"

The two answers differ for BlockExchangeSendingMailbox. Only the conservative direction kept the code correct.

Fix

Add SendingMailbox#deliversByReference(), and use it in BroadcastExchange#route:

• InMemorySendingMailbox returns true.
• GrpcSendingMailbox returns false.
• BlockExchangeSendingMailbox returns true if any mailbox of its inner exchange returns true.

The answer of a mailbox never changes, so each exchange computes it once, when it is created.

The delegation is an OR over the inner mailboxes. An inner HashExchange builds a new block for each destination, but those blocks hold the same cell objects, so one by-reference worker in a receiver stage is enough to require a copy.

isLocal() now only means what sendBlock needs. This PR also moves the contract that route depends on — a mailbox that does not deliver by reference must finish reading the block before send returns — from GrpcSendingMailbox to SendingMailbox#send, where implementers can see it.

Testing

BroadcastExchangeTest gets two spool cases: receiver stages whose workers are all remote (no copies), and a receiver stage with one worker on this server next to one on another server (one copy, shared within that stage). The first fails without this change. BlockExchangeTest covers the delegation directly. WindowFunnelTest and SpoolIntegrationTest still pass.

Out of scope

The same wrapper also reports isLocal() == true to BlockExchange#sendBlock, while the inner exchanges get BlockSplitter.NO_OP, so multi-send blocks to remote receivers are never split against MAX_MAILBOX_CONTENT_SIZE_BYTES. That defect is tracked separately in #19427.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 67.59%. Comparing base (88ea694) to head (6ca1b2f).
⚠️ Report is 67 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19428      +/-   ##
============================================
+ Coverage     67.57%   67.59%   +0.02%     
  Complexity     1430     1430              
============================================
  Files          3487     3487              
  Lines        224228   224275      +47     
  Branches      35394    35406      +12     
============================================
+ Hits         151511   151605      +94     
+ Misses        60691    60641      -50     
- Partials      12026    12029       +3     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	67.59% <100.00%> (+0.02%)	⬆️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	67.59% <100.00%> (+0.02%)	⬆️
unittests	67.59% <100.00%> (+0.02%)	⬆️
unittests1	57.69% <100.00%> (-0.01%)	⬇️
unittests2	39.34% <0.00%> (+0.02%)	⬆️

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.