Support stage identification via stage id in metrics

Author: menishmueli

spark-ui/src/interfaces/AppStore.ts

@@ -302,6 +302,7 @@ export interface EnrichedSqlNode {
   rddScopeId: string | undefined;
   icebergCommit: IcebergCommitsInfo | undefined;
   cachedStorage: RddStorageInfo | undefined;
+  nodeIdFromMetrics: number | undefined;
 }
 
 export interface SQLNodeExchangeStageData {
@@ -322,6 +323,7 @@ export interface SQLNodeStageData {
 export interface EnrichedSqlMetric {
   name: string;
   value: string;
+  stageId?: number | undefined;
 }
 
 export interface EnrichedSqlEdge {

spark-ui/src/reducers/SQLNodeStageReducer.ts

@@ -10,7 +10,7 @@ import {
 import { calculatePercentage } from "../utils/FormatUtils";
 import { generateGraph } from "./PlanGraphUtils";
 import { calculateNodeToStorageInfo } from "./SqlReducer";
-import { findExchangeStageIds, findStageIdFromMetrics, isExchangeNode } from "./SqlReducerUtils";
+import { findExchangeStageIds, isExchangeNode } from "./SqlReducerUtils";
 
 export function calculateSQLNodeStage(sql: EnrichedSparkSQL, sqlStages: SparkStagesStore): EnrichedSparkSQL {
   let nodes = sql.nodes;
@@ -316,8 +316,8 @@ export function calculateSqlStage(
       // Use exchange-specific metric parsing
       const { readStageId, writeStageId } = findExchangeStageIds(node.metrics);
       stageData = createExchangeStageData(readStageId, writeStageId, stages);
-    } else if (databricksRddStageId !== undefined || stageCodegen?.stage?.stageId !== undefined) {
-      metricsStageIdHint = findStageIdFromMetrics(node.metrics);
+    } else {
+      metricsStageIdHint = stageCodegen?.nodeIdFromMetrics ?? node.nodeIdFromMetrics;
     }
 
     stageData = stageData ?? stageDataFromStage(

spark-ui/src/reducers/SqlReducer.ts

@@ -42,6 +42,7 @@ import {
   calcNodeMetrics,
   calcNodeType,
   extractTotalFromStatisticsMetric,
+  findStageIdFromMetrics,
   nodeEnrichedNameBuilder,
 } from "./SqlReducerUtils";
 
@@ -251,7 +252,8 @@ function calculateSql(
     .filter((node) => node.isCodegenNode)
     .map((node) => {
       const codegenDuration = calcCodegenDuration(node.metrics);
-      return { ...node, codegenDuration: codegenDuration };
+      const nodeIdFromMetrics = findStageIdFromMetrics(node.metrics);
+      return { ...node, codegenDuration: codegenDuration, nodeIdFromMetrics: nodeIdFromMetrics };
     });
 
   const onlyGraphNodes = typeEnrichedNodes.filter(
@@ -276,8 +278,12 @@ function calculateSql(
       node.metrics,
     );
 
+    const nodeIdFromMetrics = findStageIdFromMetrics(node.metrics);
+    console.log(`node.metrics from sql plan: ${JSON.stringify(node.metrics)}, nodeIdFromMetrics: ${nodeIdFromMetrics}`);
+
     return {
       ...node,
+      nodeIdFromMetrics: nodeIdFromMetrics,
       metrics: updateNodeMetrics(node, node.metrics, graph, onlyGraphNodes),
       enrichedName: updateNodeEnrichedName(node, onlyGraphNodes, graph),
       parsedPlan: updateParsedPlan(node, onlyGraphNodes, graph),
@@ -449,13 +455,15 @@ export function updateSqlNodeMetrics(
 
     // TODO: cache the graph
     const graph = generateGraph(runningSql.edges, runningSql.nodes);
+    const nodeIdFromMetrics = findStageIdFromMetrics(matchedMetricsNodes[0].metrics);
     const metrics = updateNodeMetrics(node, matchedMetricsNodes[0].metrics, graph, runningSql.nodes);
     const exchangeMetrics = calcExchangeMetrics(node.nodeName, metrics);
 
     // TODO: maybe do a smarter replacement, or send only the initialized metrics
     return {
       ...node,
       metrics: metrics,
+      nodeIdFromMetrics: nodeIdFromMetrics,
       exchangeMetrics: exchangeMetrics,
     };
   });
@@ -468,12 +476,13 @@ export function updateSqlNodeMetrics(
       return node;
     }
 
+    const nodeIdFromMetrics = findStageIdFromMetrics(matchedMetricsNodes[0].metrics);
     const metrics = calcNodeMetrics(node.type, matchedMetricsNodes[0].metrics);
     const codegenDuration = calcCodegenDuration(metrics);
-
     return {
       ...node,
       codegenDuration: codegenDuration,
+      nodeIdFromMetrics: nodeIdFromMetrics,
     };
   });
   const nodesWithStorageInfo = calculateNodeToStorageInfo(stages, nodes);

spark-ui/src/reducers/SqlReducerUtils.ts

@@ -391,6 +391,12 @@ export function calcNodeMetrics(
   return (
     metrics
       .filter((metric) => allowList.includes(metric.name))
+      .map((metric) => {
+        const stageId = extractStageFromSummaryMetric(metric.value);
+        return stageId === undefined
+          ? metric
+          : { ...metric, stageId: stageId };
+      })
       .map((metric) => {
         const valueTransformer = metricsValueTransformer[metric.name];
         if (valueTransformer === undefined) {
@@ -420,7 +426,7 @@ export function calcNodeMetrics(
  * @returns The stage ID as a number, or undefined if parsing fails or no stage info is found
  */
 export function extractStageFromSummaryMetric(metricValue: string): number | undefined {
-  if (!metricValue) {
+  if (!metricValue || !metricValue.includes("(stage")) {
     return undefined;
   }
 
@@ -450,13 +456,12 @@ export function extractStageFromSummaryMetric(metricValue: string): number | und
  */
 export function findStageIdFromMetrics(metrics: EnrichedSqlMetric[]): number | undefined {
   for (const metric of metrics) {
-    if (metric.value && metric.value.includes("(stage")) {
-      const stageId = extractStageFromSummaryMetric(metric.value);
-      if (stageId !== undefined) {
-        return stageId;
-      }
+    const stageId = extractStageFromSummaryMetric(metric.value);
+    if (stageId !== undefined) {
+      return stageId;
     }
   }
+
   return undefined;
 }
 
@@ -491,15 +496,15 @@ export function isExchangeNode(nodeName: string): boolean {
 const EXCHANGE_READ_STAGE_METRICS = [
   "local bytes read",
   "fetch wait time total"
-] as const;
+];
 
 /**
  * Metric names that indicate write stage information for Exchange nodes.
  */
 const EXCHANGE_WRITE_STAGE_METRICS = [
   "shuffle write time",
   "shuffle bytes written"
-] as const;
+];
 
 /**
  * Extracts stage information from Exchange node metrics for read and write stages.
@@ -514,23 +519,14 @@ export function findExchangeStageIds(metrics: EnrichedSqlMetric[]): {
   let writeStageId: number | undefined;
 
   for (const metric of metrics) {
-    if (!metric.value || !metric.value.includes("(stage")) {
-      continue;
-    }
-
-    const stageId = extractStageFromSummaryMetric(metric.value);
-    if (stageId === undefined) {
-      continue;
-    }
-
     // Check if this metric indicates a read stage
-    if (EXCHANGE_READ_STAGE_METRICS.includes(metric.name as any)) {
-      readStageId = stageId;
+    if (EXCHANGE_READ_STAGE_METRICS.includes(metric.name)) {
+      readStageId = metric.stageId;
     }
 
     // Check if this metric indicates a write stage
-    if (EXCHANGE_WRITE_STAGE_METRICS.includes(metric.name as any)) {
-      writeStageId = stageId;
+    if (EXCHANGE_WRITE_STAGE_METRICS.includes(metric.name)) {
+      writeStageId = metric.stageId;
     }
   }