Skip block copies for spool stages whose workers are all remote

pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/GrpcSendingMailbox.java

@@ -116,9 +116,12 @@ public boolean isLocal() {
     return false;
   }
 
-  /// NOTE: [org.apache.pinot.query.runtime.operator.exchange.BlockExchange] implementations rely on this method
-  /// serializing the block synchronously on the calling thread: once it returns, the block's contents may be handed
-  /// by reference to a local receiver that mutates them.
+  @Override
+  public boolean deliversByReference() {
+    // Blocks are serialized within send(MseBlock.Data), so the receiver never reads this instance
+    return false;
+  }
+
   @Override
   public void send(MseBlock.Data data) {
     QueryThreadContext.checkTerminationAndSampleUsage(SEND_SCOPE);

pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemorySendingMailbox.java

@@ -62,6 +62,12 @@ public boolean isLocal() {
     return true;
   }
 
+  @Override
+  public boolean deliversByReference() {
+    // Blocks are offered to the receiving mailbox as they are, so the receiver reads the same instance
+    return true;
+  }
+
   @Override
   public void send(MseBlock.Data data) {
     QueryThreadContext.checkTerminationAndSampleUsage(SEND_SCOPE);

pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/SendingMailbox.java

@@ -33,14 +33,41 @@
 ///   - One call to [#cancel(Throwable)] if the sender wants to cancel the receiver
 public interface SendingMailbox extends AutoCloseable {
 
-  /// Returns whether the mailbox is sending data to a local receiver, where blocks can be directly passed to the
-  /// receiver.
+  /// Returns whether blocks can be passed to this mailbox whole, instead of being split into smaller blocks that
+  /// respect the maximum content size of a mailbox message.
+  ///
+  /// This says nothing about whether the receiver ends up with a reference to the block: see
+  /// [#deliversByReference()].
   boolean isLocal();
 
+  /// Returns whether a receiver may keep a reference to the blocks sent to this mailbox, instead of reading their
+  /// contents within [#send(MseBlock.Data)]. The answer must not change during the life of the mailbox, so callers
+  /// can read it once.
+  ///
+  /// A mailbox that returns `true` here also returns `true` from [#isLocal()], because a block can only be given by
+  /// reference to a receiver in this process. The reverse does not hold: a local mailbox can still read the block
+  /// within [#send(MseBlock.Data)] and pass on something else.
+  ///
+  /// A caller that sends the same block instance to more than one mailbox, when anything downstream mutates the
+  /// contents of that block, must:
+  ///
+  /// 1. Give a copy of the block to all the mailboxes that return `true` here, except one.
+  /// 2. Give the original block to that remaining mailbox last, once every other mailbox has returned from
+  ///    [#send(MseBlock.Data)].
+  ///
+  /// Step 2 matters as much as step 1. A receiver of the original block can mutate it as soon as `send` gives it the
+  /// block, which would corrupt the block for a mailbox that is still reading it.
+  ///
+  /// See [org.apache.pinot.query.runtime.operator.exchange.BroadcastExchange].
+  boolean deliversByReference();
+
   /// Sends a data block to the receiver. Note that SendingMailbox are required to acquire resources lazily in this
   /// call, and they should **not** acquire any resources when they are created. This method should throw if there was
   /// an error sending the data, since that would allow
   /// [org.apache.pinot.query.runtime.operator.exchange.BlockExchange] to exit early.
+  ///
+  /// Implementations that return `false` from [#deliversByReference()] must finish reading the block before this
+  /// method returns, because the caller may then pass the same block to a receiver that mutates it.
   void send(MseBlock.Data data);
 
   /// Sends an EOS block to the receiver. Note that SendingMailbox are required to acquire resources lazily in this

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MailboxSendOperator.java

@@ -89,7 +89,8 @@ public MailboxSendOperator(OpChainExecutionContext context, MultiStageOperator i
   /// 1. One inner exchange is created for each receiver stage, using the method mentioned above and keeping the
   ///    distribution type specified in the [MailboxSendNode].
   /// 2. Then, a single outer broadcast exchange is created to fan out the data to all the inner exchanges. It copies
-  ///    blocks that carry aggregation intermediate results so that no two receiver stages share them.
+  ///    blocks that carry aggregation intermediate results, so that no two receiver stages share them. Stages whose
+  ///    workers all read the block within [SendingMailbox#send(MseBlock.Data)] do not need a copy.
   ///
   /// @see BlockExchange#asSendingMailbox(String)
   private static BlockExchange getBlockExchange(OpChainExecutionContext ctx, MailboxSendNode node,

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchange.java

@@ -43,6 +43,7 @@ public abstract class BlockExchange implements AutoCloseable {
   private final List<SendingMailbox> _sendingMailboxes;
   private final BlockSplitter _splitter;
   private final Function<List<SendingMailbox>, Integer> _statsIndexChooser;
+  private final boolean _deliversByReference;
 
   protected static final Function<List<SendingMailbox>, Integer> RANDOM_INDEX_CHOOSER =
       (mailboxes) -> ThreadLocalRandom.current().nextInt(mailboxes.size());
@@ -86,6 +87,21 @@ protected BlockExchange(List<SendingMailbox> sendingMailboxes, BlockSplitter spl
     _sendingMailboxes = sendingMailboxes;
     _splitter = splitter;
     _statsIndexChooser = statsIndexChooser;
+    _deliversByReference = anyDeliversByReference(sendingMailboxes);
+  }
+
+  /// Returns whether any of the given mailboxes delivers blocks by reference. Mailboxes are fixed when the exchange
+  /// is created, and each of them gives a constant answer, so this is computed once.
+  ///
+  /// The answer therefore counts mailboxes that terminate early later on. That only makes a caller take a copy it
+  /// did not need. It never makes a caller skip a copy it did need.
+  private static boolean anyDeliversByReference(List<SendingMailbox> sendingMailboxes) {
+    for (SendingMailbox sendingMailbox : sendingMailboxes) {
+      if (sendingMailbox.deliversByReference()) {
+        return true;
+      }
+    }
+    return false;
   }
 
   /// API to send a block to the destination mailboxes.
@@ -164,6 +180,12 @@ protected void sendBlock(SendingMailbox sendingMailbox, MseBlock.Data block) {
     }
   }
 
+  /// Sends the block to the destinations, following the distribution strategy of this exchange.
+  ///
+  /// Implementations must finish reading the block before this method returns. [BlockExchangeSendingMailbox] reports
+  /// that it does not deliver blocks by reference on that basis, so an implementation that queued a block for
+  /// another thread would break [SendingMailbox#deliversByReference()] for every exchange that decorates it, and
+  /// silently corrupt the aggregation intermediate results that [BroadcastExchange] shares between destinations.
   protected abstract void route(List<SendingMailbox> destinations, MseBlock.Data block);
 
   @Override
@@ -213,9 +235,21 @@ public BlockExchangeSendingMailbox(String id) {
 
     @Override
     public boolean isLocal() {
+      // Blocks are handed to the decorated exchange whole, and splitting them is left to that exchange.
+      // TODO(#19427): the decorated exchange is currently built with BlockSplitter#NO_OP, so blocks sent
+      //       through a multi-send node are never split. See MailboxSendOperator#getBlockExchange.
       return true;
     }
 
+    @Override
+    public boolean deliversByReference() {
+      // The decorated exchange passes blocks to its own mailboxes, so this mailbox delivers by reference only if
+      // any of those does. The question is whether a receiver keeps the block alive after route returns, not whether
+      // the decorated exchange shares one block between its own mailboxes: an inner HashExchange gives each of its
+      // mailboxes a different block, but the rows of those blocks still hold the cells of this one.
+      return _deliversByReference;
+    }
+
     @Override
     public void send(MseBlock.Data data) {
       if (LOGGER.isTraceEnabled()) {

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/exchange/BroadcastExchange.java

@@ -28,14 +28,14 @@
 
 /// Broadcast blocks to all the destinations.
 ///
-/// This is the only exchange that routes the same block to more than one destination, and local (same-JVM) mailboxes
-/// deliver on-heap blocks by reference. Blocks that
+/// This is the only exchange that routes the same block to more than one destination, and some mailboxes
+/// [deliver blocks by reference][SendingMailbox#deliversByReference()]. Blocks that
 /// [carry aggregation intermediate results][RowHeapDataBlock#containsObjectColumns()] cannot be shared this way:
 /// downstream operators mutate those objects in place when they merge them or extract final results, so two
-/// receivers on the same server would corrupt them. For such blocks, [#route] gives every local destination except
-/// the first its own [copy][RowHeapDataBlock#copyObjectColumns()]. Remote destinations only read the block to
-/// serialize it, before any local receiver can mutate it, so they do not need copies. Blocks without OBJECT columns
-/// are shared by reference with all the destinations.
+/// receivers of the same block would corrupt them. For such blocks, [#route] gives every destination that delivers
+/// by reference, except the first, its own [copy][RowHeapDataBlock#copyObjectColumns()]. The other destinations read
+/// the block within [SendingMailbox#send(MseBlock.Data)], before any receiver can mutate it, so they do not need
+/// copies. Blocks without OBJECT columns are shared by reference with all the destinations.
 ///
 /// This also protects multi-send (spool) nodes, which fan each block out to the exchanges of their receiver stages
 /// through this exchange (see [BlockExchange#asSendingMailbox]).
@@ -60,26 +60,26 @@ protected void route(List<SendingMailbox> destinations, MseBlock.Data block) {
       }
       return;
     }
-    // Send a copy to every active local destination except the first one, which receives the original block without
-    // copying. Remote destinations serialize the original block on this thread, and the copies are also made on this
-    // thread, so all reads of the original block finish before it is handed to a local receiver that can start
-    // mutating it.
+    // Send a copy to every active destination that delivers by reference, except the first one, which receives the
+    // original block without copying. The other destinations read the original block within send, and the copies are
+    // also made on this thread, so all reads of the original block finish before it is handed to a receiver that can
+    // start mutating it.
     RowHeapDataBlock rowHeapBlock = block.asRowHeap();
-    SendingMailbox firstLocalDestination = null;
+    SendingMailbox firstByReferenceDestination = null;
     for (SendingMailbox mailbox : destinations) {
       if (mailbox.isEarlyTerminated()) {
         continue;
       }
-      if (!mailbox.isLocal()) {
+      if (!mailbox.deliversByReference()) {
         sendBlock(mailbox, block);
-      } else if (firstLocalDestination == null) {
-        firstLocalDestination = mailbox;
+      } else if (firstByReferenceDestination == null) {
+        firstByReferenceDestination = mailbox;
       } else {
         sendBlock(mailbox, rowHeapBlock.copyObjectColumns());
       }
     }
-    if (firstLocalDestination != null) {
-      sendBlock(firstLocalDestination, block);
+    if (firstByReferenceDestination != null) {
+      sendBlock(firstByReferenceDestination, block);
     }
   }
 }

pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/GrpcSendingMailboxTest.java

@@ -80,6 +80,18 @@ public void sendDataThrowsWhenQueryTerminated() {
     }
   }
 
+  @Test
+  public void doesNotDeliverBlocksByReference() {
+    GrpcSendingMailbox mailbox =
+        new GrpcSendingMailbox("test-mailbox", Mockito.mock(ChannelManager.class), "localhost", 0, Long.MAX_VALUE,
+            new StatMap<>(MailboxSendOperator.StatKey.class), 4 * 1024 * 1024, true);
+
+    // Blocks are serialized within send(MseBlock.Data), so senders that share one block between mailboxes do not
+    // need to give this one a copy. See BroadcastExchange.
+    Assert.assertFalse(mailbox.deliversByReference());
+    Assert.assertFalse(mailbox.isLocal());
+  }
+
   /// Regression test for the lazy-initialization data race on `_contentObserver`. Before the fix, both `sendInternal`
   /// and `cancel` had an unsynchronized `if (_contentObserver == null) { _contentObserver = getContentObserver(); }`
   /// pattern. `_contentObserver` is `volatile` so individual reads/writes are atomic, but two threads racing through

pinot-query-runtime/src/test/java/org/apache/pinot/query/mailbox/InMemorySendingMailboxTest.java

@@ -49,4 +49,16 @@ public void sendDataThrowsWhenQueryTerminated() {
       Mockito.verifyNoInteractions(mailboxService);
     }
   }
+
+  @Test
+  public void deliversBlocksByReference() {
+    InMemorySendingMailbox mailbox =
+        new InMemorySendingMailbox("test-mailbox", Mockito.mock(MailboxService.class), Long.MAX_VALUE,
+            new StatMap<>(MailboxSendOperator.StatKey.class));
+
+    // Blocks are offered to the receiving mailbox as they are, so senders that share one block between mailboxes
+    // must give this one a copy. See BroadcastExchange.
+    Assert.assertTrue(mailbox.deliversByReference());
+    Assert.assertTrue(mailbox.isLocal());
+  }
 }

pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/exchange/BlockExchangeTest.java

@@ -165,6 +165,28 @@ public void shouldSplitBlocks()
     Assert.assertEquals(sentBlocks.get(1).asRowHeap().getRows(), outBlockTwo.getRows());
   }
 
+  @Test
+  public void shouldDeliverByReferenceWhenAnyDestinationDoes() {
+    // Given: an exchange whose destinations are one mailbox that serializes blocks and one that does not
+    when(_mailbox2.deliversByReference()).thenReturn(true);
+    BlockExchange exchange = new TestBlockExchange(List.of(_mailbox1, _mailbox2));
+
+    // Then: the exchange exposed as a mailbox reports that it delivers by reference
+    Assert.assertTrue(exchange.asSendingMailbox("1").deliversByReference());
+  }
+
+  @Test
+  public void shouldNotDeliverByReferenceWhenNoDestinationDoes() {
+    // Given: an exchange whose destinations all serialize the blocks they are sent
+    BlockExchange exchange = new TestBlockExchange(List.of(_mailbox1, _mailbox2));
+
+    // Then: the exchange exposed as a mailbox reports that it does not deliver by reference. It still takes blocks
+    // whole, because splitting them is left to the exchange it decorates
+    SendingMailbox sendingMailbox = exchange.asSendingMailbox("1");
+    Assert.assertFalse(sendingMailbox.deliversByReference());
+    Assert.assertTrue(sendingMailbox.isLocal());
+  }
+
   private static class TestBlockExchange extends BlockExchange {
     protected TestBlockExchange(List<SendingMailbox> destinations) {
       this(destinations, (block, size) -> Iterators.singletonIterator(block));