gortiz commented on code in PR #19816:
URL: https://github.com/apache/pinot/pull/19816#discussion_r4273757119
##########
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:
I checked `InStageStatsTreeBuilder` and I agree with Yash for the legacy
mode. `getPipelineBreakerNode()` checks that the entry before the current index
is the `PIPELINE_BREAKER`, then `_index--` moves to it and
`visitMailboxReceive()` calls `selfNode()`, which copies the breaker's stat
map. So the node under `LEAF` carries the breaker's own stats, and the equality
holds with `streamStats = false`.
That said, the helper now relies on this rendering detail (the breaker's
stats under a `MAILBOX_RECEIVE` label), and it also assumes a single pipeline
breaker per stage, like the renderer does. Could we also add a small unit test
over raw `StageStats.Closed` lists, with distinct allocations for the breaker,
its receive, the leaf threads and the stage root? It would cover both the
legacy shape and the flattened stream shape, it would not break if someone
fixes the renderer, and it runs in CI even before #19856 is merged. The
integration test can stay 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]
