gortiz commented on code in PR #12704: URL: https://github.com/apache/pinot/pull/12704#discussion_r1576190502
########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java: ########## @@ -276,6 +264,161 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S return brokerResponse; } + private void fillOldBrokerResponseStats(BrokerResponseNativeV2 brokerResponse, + List<MultiStageQueryStats.StageStats.Closed> queryStats, DispatchableSubPlan dispatchableSubPlan) { + for (int i = 0; i < queryStats.size(); i++) { + MultiStageQueryStats.StageStats.Closed stageStats = queryStats.get(i); + if (stageStats == null) { + brokerResponse.addStageStats(JsonUtils.newObjectNode()); + } else { + stageStats.forEach((type, stats) -> type.mergeInto(brokerResponse, stats)); + + DispatchablePlanFragment dispatchablePlanFragment = dispatchableSubPlan.getQueryStageList().get(i); + MultiStageStatsTreeBuilder treeBuilder = new MultiStageStatsTreeBuilder(stageStats); + PlanNode fragmentRoot = dispatchablePlanFragment.getPlanFragment().getFragmentRoot(); + JsonNode node = fragmentRoot.visit(treeBuilder, null); + brokerResponse.addStageStats(node); + } + } + } + + public static class MultiStageStatsTreeBuilder implements PlanNodeVisitor<JsonNode, Void> { + private final MultiStageQueryStats.StageStats.Closed _stageStats; + private int _index; + private static final String CHILDREN_KEY = "children"; + + public MultiStageStatsTreeBuilder(MultiStageQueryStats.StageStats.Closed stageStats) { + _stageStats = stageStats; + _index = stageStats.getLastOperatorIndex(); + } + + private ObjectNode selfNode(MultiStageOperator.Type type) { + ObjectNode json = JsonUtils.newObjectNode(); + json.put("type", type.toString()); + Iterator<Map.Entry<String, JsonNode>> statsIt = _stageStats.getOperatorStats(_index).asJson().fields(); + while (statsIt.hasNext()) { + Map.Entry<String, JsonNode> entry = statsIt.next(); + json.set(entry.getKey(), entry.getValue()); + } + return json; + } + + private JsonNode recursiveCase(AbstractPlanNode node, MultiStageOperator.Type expectedType) { + MultiStageOperator.Type type = _stageStats.getOperatorType(_index); + if (type != expectedType) { Review Comment: No, this is very rare. Mostly only happens on the stage 0 and in leaf stages. In stage 0 always have two nodes. One is a transform IIRC and the other is a receive node and we only have stats for the latter. In leaf stage, we may have several nodes, but the leaf stage compiles them into a different plan (as explained above). In can add some comments -- 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: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org