yashmayya commented on code in PR #18736:
URL: https://github.com/apache/pinot/pull/18736#discussion_r3416482164


##########
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/MultiStageStatsTreeBuilder.java:
##########
@@ -63,4 +96,135 @@ public ObjectNode jsonStatsByStage(int stage) {
     InStageStatsTreeBuilder treeBuilder = new 
InStageStatsTreeBuilder(stageStats, this::jsonStatsByStage);
     return planNode.visit(treeBuilder, new InStageStatsTreeBuilder.Context(1));
   }
+
+  /**
+   * Renders a stage's explicit stats tree. The shape (and any plugin operator 
types) come straight from the decoded
+   * tree; the stage's PlanNode tree is consulted only to resolve which plan 
nodes are {@link MailboxReceiveNode}s so
+   * the sender stage's tree can be nested under them, mirroring the legacy 
renderer's cross-stage nesting.
+   */
+  private ObjectNode jsonFromStatsTree(StageStatsTreeNode statsTree, int 
stage) {
+    Map<Integer, PlanNode> planNodesById = new HashMap<>();
+    PlanNode fragmentRoot = _planNodes.get(stage);
+    if (fragmentRoot != null) {
+      assignPlanNodeIds(fragmentRoot, planNodesById, new int[]{0});
+    }
+    // The stage root is the send operator: its parallelism (when reported, 
e.g. by the built-in MAILBOX_SEND stats)
+    // scales cpu-time into wall-clock time for the whole stage. Plugin send 
types may not report it; default to 1.
+    int parallelism = 1;
+    JsonNode rootParallelism = 
statsTree.getStatMap().asJson().get("parallelism");
+    if (rootParallelism != null) {
+      parallelism = Math.max(1, rootParallelism.asInt(1));
+    }
+    return jsonFromStatsTreeNode(statsTree, planNodesById, parallelism);
+  }
+
+  private ObjectNode jsonFromStatsTreeNode(StageStatsTreeNode node, 
Map<Integer, PlanNode> planNodesById,
+      int parallelism) {
+    ObjectNode json = JsonUtils.newObjectNode();
+    json.put("type", node.getType().name());
+    for (Map.Entry<String, JsonNode> entry : 
node.getStatMap().asJson().properties()) {
+      json.set(entry.getKey(), entry.getValue());
+    }
+    if (json.get("parallelism") == null) {
+      json.put("parallelism", parallelism);
+    }
+    JsonNode executionTimeMs = json.get("executionTimeMs");
+    if (executionTimeMs != null) {
+      json.put("clockTimeMs", executionTimeMs.asLong(0) / parallelism);
+    }
+
+    if (!node.getPlanNodeIds().isEmpty()) {
+      ArrayNode planNodeIds = JsonUtils.newArrayNode();
+      node.getPlanNodeIds().forEach(planNodeIds::add);
+      json.set("planNodeIds", planNodeIds);
+    }
+
+    ArrayNode children = JsonUtils.newArrayNode();
+    for (StageStatsTreeNode child : node.getChildren()) {
+      // Collapse PIPELINE_BREAKER nodes for parity with the legacy renderer, 
which nests the
+      // breaker's receive operators directly under the LEAF and never renders 
the breaker itself.
+      // The JSON shape is consumed by external tooling, so the explicit-tree 
mode must not change it.
+      if (child.getType().getId() == 
MultiStageOperator.Type.PIPELINE_BREAKER.getId()) {
+        for (StageStatsTreeNode grandChild : child.getChildren()) {
+          children.add(jsonFromStatsTreeNode(grandChild, planNodesById, 
parallelism));
+        }
+      } else {
+        children.add(jsonFromStatsTreeNode(child, planNodesById, parallelism));
+      }
+    }
+    // Cross-stage nesting: a node whose plan node is a mailbox receive gets 
the sender stage's tree as a child.
+    // Type-agnostic on purpose — plugin-defined receive operators carry a 
different descriptor than the built-in
+    // MAILBOX_RECEIVE but map to the same MailboxReceiveNode.
+    for (Integer planNodeId : node.getPlanNodeIds()) {
+      PlanNode planNode = planNodesById.get(planNodeId);
+      if (planNode instanceof MailboxReceiveNode) {
+        children.add(jsonStatsByStage(((MailboxReceiveNode) 
planNode).getSenderStageId()));
+      }
+    }
+    if (!children.isEmpty()) {
+      json.set("children", children);
+    }
+
+    // self* fields: parent-minus-children derived stats, matching 
InStageStatsTreeBuilder semantics.
+    // These must be omitted when zero to keep parity with the legacy renderer.
+    JsonNode execNode = json.get("executionTimeMs");
+    if (execNode != null) {
+      long selfExecTimeMs = execNode.asLong(0) - sumChildrenStat(node, 
"executionTimeMs");
+      if (selfExecTimeMs != 0) {
+        json.put("selfExecutionTimeMs", selfExecTimeMs);
+        json.put("selfClockTimeMs", selfExecTimeMs / parallelism);
+      }
+    }
+    JsonNode allocNode = json.get("allocatedMemoryBytes");
+    if (allocNode != null) {
+      long selfAllocBytes = allocNode.asLong(0) - sumChildrenStat(node, 
"allocatedMemoryBytes");
+      if (selfAllocBytes != 0) {
+        json.put("selfAllocatedMB", selfAllocBytes / (1024 * 1024));
+      }
+    }
+    JsonNode gcNode = json.get("gcTimeMs");
+    if (gcNode != null) {
+      long selfGcTimeMs = gcNode.asLong(0) - sumChildrenStat(node, "gcTimeMs");
+      if (selfGcTimeMs != 0) {
+        json.put("selfGcTimeMs", selfGcTimeMs);
+      }
+    }
+
+    return json;
+  }
+
+  /**
+   * Sums a named stat field across the logical children of a node (with 
PIPELINE_BREAKER collapse,
+   * mirroring the child iteration in {@link #jsonFromStatsTreeNode}).
+   */
+  private static long sumChildrenStat(StageStatsTreeNode node, String statKey) 
{
+    long sum = 0;
+    for (StageStatsTreeNode child : node.getChildren()) {
+      if (child.getType().getId() == 
MultiStageOperator.Type.PIPELINE_BREAKER.getId()) {

Review Comment:
   Minor / non-blocking (parity nit): `sumChildrenStat` collapses the 
`PIPELINE_BREAKER` and subtracts its receive's stats from the LEAF — but the 
legacy renderer deliberately does **not** subtract the breaker. 
`InStageStatsTreeBuilder` renders a PB-bearing LEAF via `selfNode(..., 
adjustWithChildren=false)`, so the leaf's `self*` stays its full value (the 
breaker ran pre-stage, so its time isn't part of the leaf's cumulative time).
   
   Net effect: for a semi-join / lookup-join leaf, `selfExecutionTimeMs` / 
`selfClockTimeMs` / `selfAllocatedMB` / `selfGcTimeMs` come out smaller than 
the legacy renderer by the breaker's contribution. I confirmed it with a unit 
test — a `LEAF` with `executionTimeMs=100` over a PB whose receive has 
`executionTimeMs=30` renders `selfExecutionTimeMs=70`, where legacy gives `100`.
   
   Small fix: skip `PIPELINE_BREAKER` children in `sumChildrenStat` (don't 
recurse into the grandchildren) to mirror `adjustWithChildren=false`. 
Everything else in the `self*` computation matches legacy (including correctly 
*not* subtracting the cross-stage sender at a `MAILBOX_RECEIVE`). 
Diagnostic-field only, so not a blocker — flagging for a follow-up since it's 
the one case this renderer was extended for.



-- 
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