Jackie-Jiang commented on code in PR #19816:
URL: https://github.com/apache/pinot/pull/19816#discussion_r4264690461


##########
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"));

Review Comment:
   The child rendered under a `LEAF` is the breaker's `MAILBOX_RECEIVE`, not 
the `PIPELINE_BREAKER` itself: the stats renderer omits the breaker node. 
Production adds the breaker's allocation, which also includes work outside 
receive calls, such as growing `dataBlocks` in 
`PipelineBreakerOperator.getNextBlock()`. This helper therefore misses those 
allocations for the semi-join, so the exact equality at line 340 is not valid 
once the test runs. Could we verify that contribution using raw stage stats 
with distinct breaker/receive allocations, or expose the breaker allocation 
before deriving the expected total from JSON?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to