Add totalMemAllocatedBytes to multi-stage query responses

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

@@ -1022,6 +1022,11 @@ private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse,
       for (MultiStageQueryStats.StageStats.Closed stageStats : queryStats) {
         if (stageStats != null) { // for example pipeline breaker may not have stats
           stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats));
+          if (!stageStats.isEmpty()) {
+            // The stage root's allocation covers all the operators of its op chain
+            brokerResponse.addTotalMemAllocatedBytes(
+                stageStats.getLastOperatorStats().getUnsafe("ALLOCATED_MEMORY_BYTES", 0L));
+          }
         }
       }
       // Broker-pruned segments are computed during routing in WorkerManager, not reported by servers.

pinot-common/src/main/java/org/apache/pinot/common/response/broker/BrokerResponseNativeV2.java

@@ -53,9 +53,10 @@
     "realtimeResponseSerializationCpuTimeNs", "offlineTotalCpuTimeNs", "realtimeTotalCpuTimeNs",
     "explainPlanNumEmptyFilterSegments", "explainPlanNumMatchAllFilterSegments", "traceInfo", "tablesQueried",
     "offlineThreadMemAllocatedBytes", "realtimeThreadMemAllocatedBytes", "offlineResponseSerMemAllocatedBytes",
-    "realtimeResponseSerMemAllocatedBytes", "offlineTotalMemAllocatedBytes", "realtimeTotalMemAllocatedBytes", "pools",
-    "rlsFiltersApplied", "approximateFunctionApplied", "groupsTrimmed", "mseLiteLeafStageLimitReached",
-    "mseLiteLeafStageEffectiveLimit", "mseLiteFanOutAdjustedLimitApplied", "responseMetadata"
+    "realtimeResponseSerMemAllocatedBytes", "offlineTotalMemAllocatedBytes", "realtimeTotalMemAllocatedBytes",
+    "totalMemAllocatedBytes", "pools", "rlsFiltersApplied", "approximateFunctionApplied", "groupsTrimmed",
+    "mseLiteLeafStageLimitReached", "mseLiteLeafStageEffectiveLimit", "mseLiteFanOutAdjustedLimitApplied",
+    "responseMetadata"
 })
 public class BrokerResponseNativeV2 implements BrokerResponse {
   private final StatMap<StatKey> _brokerStats = new StatMap<>(StatKey.class);
@@ -81,6 +82,7 @@ public class BrokerResponseNativeV2 implements BrokerResponse {
   /// virtual rows can be generated. For example, in a join query, the number of rows can be more than the number of
   /// rows in the table.
   private long _maxRowsInOperator;
+  private long _totalMemAllocatedBytes;
   private String _requestId;
   private String _clientRequestId;
   private String _brokerId;
@@ -272,6 +274,17 @@ public void mergeMaxRowsInOperator(long maxRows) {
     _maxRowsInOperator = Math.max(_maxRowsInOperator, maxRows);
   }
 
+  /// Returns the heap bytes allocated by the query across all its stages and workers, as reported in the stage stats.
+  /// This is the total allocated (including short-lived objects), not the peak heap used. It is 0 unless thread
+  /// allocated bytes measurement is enabled on the brokers and servers.
+  public long getTotalMemAllocatedBytes() {
+    return _totalMemAllocatedBytes;
+  }
+
+  public void addTotalMemAllocatedBytes(long bytes) {
+    _totalMemAllocatedBytes += bytes;
+  }
+
   @Override
   public long getTimeUsedMs() {
     return _timeUsedMs;

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiNodesOfflineClusterIntegrationTest.java

@@ -326,6 +326,35 @@ public void testConstantExpressionQuery()
     assertEquals(result.get("numServersQueried").intValue(), 1);
   }
 
+  @Test
+  public void testMultiStageTotalMemAllocatedBytes()
+      throws Exception {
+    setUseMultiStageQueryEngine(true);
+
+    // The semi-join query uses a pipeline breaker
+    for (String query : List.of("SELECT Carrier, COUNT(*) FROM mytable GROUP BY Carrier",
+        "SELECT Carrier FROM mytable WHERE Origin IN (SELECT Origin FROM mytable WHERE ArrDelay > 0) LIMIT 10")) {
+      JsonNode result = postQuery(query);
+      long totalMemAllocatedBytes = result.path("totalMemAllocatedBytes").asLong();
+      assertTrue(totalMemAllocatedBytes > 0);
+      assertEquals(totalMemAllocatedBytes, sumAllocatedBytes(result.get("stageStats"), true));
+    }
+  }
+
+  /// Sums the allocation of each op chain root (the top node, each MAILBOX_SEND and each pipeline breaker, which is
+  /// rendered as the child of its LEAF) and of the single-stage threads of each LEAF.
+  private static long sumAllocatedBytes(JsonNode node, boolean isOpChainRoot) {
+    long bytes = node.path("threadMemAllocatedBytes").asLong();
+    if (isOpChainRoot) {
+      bytes += node.path("allocatedMemoryBytes").asLong();
+    }
+    boolean isLeaf = node.path("type").asText().equals("LEAF");
+    for (JsonNode child : node.path("children")) {
+      bytes += sumAllocatedBytes(child, isLeaf || child.path("type").asText().equals("MAILBOX_SEND"));
+    }
+    return bytes;
+  }
+
   // Disabled because segments might not be server partitioned with multiple servers
   @Test(enabled = false)
   @Override

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

@@ -422,6 +422,8 @@ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
         @SuppressWarnings("unchecked")
         StatMap<LeafOperator.StatKey> stats = (StatMap<LeafOperator.StatKey>) map;
         response.mergeMaxRowsInOperator(stats.getLong(LeafOperator.StatKey.EMITTED_ROWS));
+        // The single-stage threads of the leaf are not covered by the stage root's allocation
+        response.addTotalMemAllocatedBytes(stats.getLong(LeafOperator.StatKey.THREAD_MEM_ALLOCATED_BYTES));
 
         StatMap<BrokerResponseNativeV2.StatKey> brokerStats = new StatMap<>(BrokerResponseNativeV2.StatKey.class);
         for (LeafOperator.StatKey statKey : stats.keySet()) {
@@ -495,6 +497,8 @@ public void mergeInto(BrokerResponseNativeV2 response, StatMap<?> map) {
         @SuppressWarnings("unchecked")
         StatMap<PipelineBreakerOperator.StatKey> stats = (StatMap<PipelineBreakerOperator.StatKey>) map;
         response.mergeMaxRowsInOperator(stats.getLong(PipelineBreakerOperator.StatKey.EMITTED_ROWS));
+        // Pipeline breakers run in their own op chain, so they are not covered by the stage root's allocation
+        response.addTotalMemAllocatedBytes(stats.getLong(PipelineBreakerOperator.StatKey.ALLOCATED_MEMORY_BYTES));
       }
     },
     SORT_OR_LIMIT(10, SortOperator.StatKey.class) {