yashmayya commented on code in PR #19816:
URL: https://github.com/apache/pinot/pull/19816#discussion_r4275999471
##########
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:
Thanks @gortiz, good idea. I added
`MultiStageBrokerRequestHandlerTest#testMergeStageStatsTotalMemAllocatedBytes`
in 60c80d250f.
- The per-stage loop moved without changes from
`fillOldBrokerResponseStats()` to `mergeStageStats()`, so the test runs the
production code over raw `StageStats.Closed` lists.
- Each operator has a distinct allocation: root receive 100000, send 10000,
leaf threads 1000, leaf op chain 100, breaker 10, breaker receive 1. The
expected total is 111010.
- So the test fails if the breaker is missing (111000), if the stage-root
lookup breaks (1010), or if an operator that its root already covers is counted
again. I checked the first two cases locally.
- The flat list for a leaf stage with a breaker has the same order in legacy
and stream modes, so the test does not depend on the JSON rendering.
The integration test stays as a smoke test.
--
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]
