Copy blocks with mutable agg intermediates per local receiver in MSE broadcast exchange

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/WindowFunnelTest.java

@@ -21,12 +21,15 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import java.io.File;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import org.apache.pinot.integration.tests.window.utils.WindowFunnelUtils;
 import org.apache.pinot.spi.data.Schema;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
 
 
 @Test(suiteName = "CustomClusterIntegrationTest")
@@ -1075,6 +1078,75 @@ public void testFunnelStepDurationStatsGroupByQueries3(boolean useMultiStageQuer
     assertEquals(row.get(9).doubleValue(), 10.0);
   }
 
+  /// Regression test for a spooled stage whose output contains funnel intermediate results (priority queues of step
+  /// events): the same stage (scan + partial funnelStepDurationStats aggregation) feeds two different consumer
+  /// stages. Without copying the blocks per receiver stage, the consumers on the same server share (and corrupt) the
+  /// same mutable priority queues when merging them or extracting final results.
+  @Test(dataProvider = "useV2QueryEngine")
+  public void testFunnelStepDurationStatsGroupByQueriesWithSpools(boolean useMultiStageQueryEngine)
+      throws Exception {
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+    // The spooled partial aggregation runs in the leaf stage
+    checkSpooledFunnelStepDurationStats(String.format("FROM %s", getTableName()));
+    // The spooled partial aggregation runs in an intermediate stage, on top of a window function CTE
+    checkSpooledFunnelStepDurationStats(String.format(
+        "FROM (SELECT userId, timestampCol, url, ROW_NUMBER() OVER (PARTITION BY url ORDER BY timestampCol) AS occ "
+            + "FROM %s) WHERE occ >= 1", getTableName()));
+  }
+
+  private void checkSpooledFunnelStepDurationStats(String fromClause)
+      throws Exception {
+    String query =
+        "SET useSpools = true; "
+            + "WITH durationStats AS (SELECT "
+            + "userId, funnelStepDurationStats(timestampCol, '1000', 4, "
+            + "url = '/product/search', "
+            + "url = '/cart/add', "
+            + "url = '/checkout/start', "
+            + "url = '/checkout/confirmation', "
+            + "'durationFunctions=count,avg,median' "
+            + ") as stats "
+            + fromClause + " GROUP BY userId) "
+            + "SELECT * FROM "
+            + "(SELECT SUM(arrayElementAtDouble(stats, 1)) AS totalCount FROM durationStats) "
+            + "CROSS JOIN "
+            + "(SELECT AVG(arrayElementAtDouble(stats, 2)) AS avgAvgDuration FROM durationStats)";
+    JsonNode jsonNode = postQuery(query);
+    assertNoError(jsonNode);
+    assertSpooled(jsonNode);
+    JsonNode rows = jsonNode.get("resultTable").get("rows");
+    assertEquals(rows.size(), 1);
+    JsonNode row = rows.get(0);
+    assertEquals(row.size(), 2);
+    assertEquals(row.get(0).doubleValue(), 40.0);
+    assertEquals(row.get(1).doubleValue(), 2.5);
+  }
+
+  /// Asserts the query actually used a spool: a spooled stage's stats subtree is rendered under each of its receiver
+  /// stages, so its MAILBOX_SEND stage id occurs more than once in the stage stats tree. Without a spool every stage
+  /// has a single parent and occurs exactly once. Guards the spool regression tests against planner changes that stop
+  /// deduplicating the shared subtree (which would make them pass vacuously).
+  private static void assertSpooled(JsonNode jsonNode) {
+    Map<Integer, Integer> sendStageCounts = new HashMap<>();
+    countMailboxSendStages(jsonNode.get("stageStats"), sendStageCounts);
+    assertTrue(sendStageCounts.values().stream().anyMatch(count -> count > 1),
+        "Expected a spooled stage (same MAILBOX_SEND stage under multiple parents), got stage counts: "
+            + sendStageCounts);
+  }
+
+  private static void countMailboxSendStages(JsonNode node, Map<Integer, Integer> sendStageCounts) {
+    if (node == null || (!node.isObject() && !node.isArray())) {
+      return;
+    }
+    if (node.isObject() && node.has("type") && "MAILBOX_SEND".equals(node.get("type").asText())
+        && node.has("stage")) {
+      sendStageCounts.merge(node.get("stage").asInt(), 1, Integer::sum);
+    }
+    for (JsonNode child : node) {
+      countMailboxSendStages(child, sendStageCounts);
+    }
+  }
+
   @Override
   public String getTableName() {
     return WindowFunnelUtils.DEFAULT_TABLE_NAME;

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

@@ -116,6 +116,9 @@ 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 void send(MseBlock.Data data) {
     QueryThreadContext.checkTerminationAndSampleUsage(SEND_SCOPE);

pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/blocks/RowHeapDataBlock.java

@@ -18,13 +18,19 @@
  */
 package org.apache.pinot.query.runtime.blocks;
 
+import com.google.common.base.Preconditions;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.List;
 import javax.annotation.Nullable;
+import org.apache.pinot.common.CustomObject;
 import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
 import org.apache.pinot.core.common.datablock.DataBlockBuilder;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.spi.query.QueryThreadContext;
 
 
 /// A block that contains data in row heap format.
@@ -102,6 +108,57 @@ public RowHeapDataBlock asRowHeap() {
     return this;
   }
 
+  /// Returns whether this block contains [OBJECT][ColumnDataType#OBJECT] columns, which hold mutable aggregation
+  /// intermediate results — the only cell values downstream operators mutate in place. Blocks that contain them must
+  /// not be shared by reference across receivers; see [#copyObjectColumns()]. Extend both methods if operators ever
+  /// start to mutate other cell types.
+  public boolean containsObjectColumns() {
+    for (ColumnDataType storedType : _dataSchema.getStoredColumnDataTypes()) {
+      if (storedType == ColumnDataType.OBJECT) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /// Returns a copy of this block whose [OBJECT][ColumnDataType#OBJECT] cells do not share state with this block.
+  ///
+  /// OBJECT cells hold aggregation intermediate results (e.g. priority queues or sketches) that downstream operators
+  /// mutate in place when they merge them or extract final results, so the same instance must not reach two
+  /// receivers. This method copies them by round-tripping them through the corresponding aggregation function's
+  /// intermediate result serde, which requires the aggregation functions to be attached to this block whenever a
+  /// non-null OBJECT cell is present (the same requirement [DataBlockBuilder] imposes to serialize such blocks).
+  ///
+  /// The row arrays are cloned, but cells of all other column types are shared with the returned block, even though
+  /// some of them are backed by mutable objects (maps, arrays, byte arrays): downstream operators do not mutate
+  /// them. Extend this method if that ever changes.
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  public RowHeapDataBlock copyObjectColumns() {
+    ColumnDataType[] storedTypes = _dataSchema.getStoredColumnDataTypes();
+    int numColumns = storedTypes.length;
+    List<Object[]> copiedRows = new ArrayList<>(_rows.size());
+    int numCellsProcessed = 0;
+    for (Object[] row : _rows) {
+      Object[] copiedRow = row.clone();
+      for (int colId = 0; colId < numColumns; colId++) {
+        Object value = copiedRow[colId];
+        if (value != null && storedTypes[colId] == ColumnDataType.OBJECT) {
+          Preconditions.checkState(_aggFunctions != null,
+              "Cannot copy OBJECT column: %s without aggregation functions", _dataSchema.getColumnName(colId));
+          QueryThreadContext.checkTerminationAndSampleUsagePeriodically(numCellsProcessed++,
+              "RowHeapDataBlock#copyObjectColumns");
+          // NOTE: The first (numColumns - numAggFunctions) columns are key columns
+          AggregationFunction aggFunction = _aggFunctions[colId + _aggFunctions.length - numColumns];
+          AggregationFunction.SerializedIntermediateResult serialized = aggFunction.serializeIntermediateResult(value);
+          copiedRow[colId] = aggFunction.deserializeIntermediateResult(
+              new CustomObject(serialized.getType(), ByteBuffer.wrap(serialized.getBytes())));
+        }
+      }
+      copiedRows.add(copiedRow);
+    }
+    return new RowHeapDataBlock(copiedRows, _dataSchema, _aggFunctions);
+  }
+
   @Override
   public SerializedDataBlock asSerialized() {
     try {

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

@@ -88,7 +88,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.
+  /// 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.
   ///
   /// @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/BroadcastExchange.java

@@ -23,9 +23,22 @@
 import org.apache.pinot.query.mailbox.SendingMailbox;
 import org.apache.pinot.query.runtime.blocks.BlockSplitter;
 import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock;
 
 
-/// Broadcast blocks to all receiving servers.
+/// 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
+/// [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.
+///
+/// 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]).
 class BroadcastExchange extends BlockExchange {
 
   protected BroadcastExchange(List<SendingMailbox> sendingMailboxes, BlockSplitter splitter) {
@@ -39,8 +52,34 @@ protected BroadcastExchange(List<SendingMailbox> sendingMailboxes, BlockSplitter
 
   @Override
   protected void route(List<SendingMailbox> destinations, MseBlock.Data block) {
+    // Serialized blocks are read-only (every receiver deserializes its own copy of the data), so they are always
+    // safe to share
+    if (destinations.size() == 1 || !block.isRowHeap() || !block.asRowHeap().containsObjectColumns()) {
+      for (SendingMailbox mailbox : destinations) {
+        sendBlock(mailbox, 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.
+    RowHeapDataBlock rowHeapBlock = block.asRowHeap();
+    SendingMailbox firstLocalDestination = null;
     for (SendingMailbox mailbox : destinations) {
-      sendBlock(mailbox, block);
+      if (mailbox.isEarlyTerminated()) {
+        continue;
+      }
+      if (!mailbox.isLocal()) {
+        sendBlock(mailbox, block);
+      } else if (firstLocalDestination == null) {
+        firstLocalDestination = mailbox;
+      } else {
+        sendBlock(mailbox, rowHeapBlock.copyObjectColumns());
+      }
+    }
+    if (firstLocalDestination != null) {
+      sendBlock(firstLocalDestination, block);
     }
   }
 }