group summary nodes by stage id

spark-ui/src/components/SqlFlow/SqlFlow.tsx

@@ -1,24 +1,41 @@
 import React, { FC, useCallback, useEffect, useState } from "react";
 import ReactFlow, {
+  addEdge,
   ConnectionLineType,
   Controls,
   ReactFlowInstance,
-  addEdge,
   useEdgesState,
   useNodesState,
 } from "reactflow";
 
-import { Box, Drawer, FormControl, InputLabel, MenuItem, Select } from "@mui/material";
+import {
+  Box,
+  Drawer,
+  FormControl,
+  InputLabel,
+  MenuItem,
+  Select,
+} from "@mui/material";
 import "reactflow/dist/style.css";
 import { useAppDispatch, useAppSelector } from "../../Hooks";
 import { EnrichedSparkSQL, GraphFilter } from "../../interfaces/AppStore";
-import { setSQLMode, setSelectedStage } from "../../reducers/GeneralSlice";
-import SqlLayoutService from "./SqlLayoutService";
-import StageIconDrawer from "./StageIconDrawer";
-import { StageNode, StageNodeName } from "./StageNode";
+import { setSelectedStage, setSQLMode } from "../../reducers/GeneralSlice";
+import StageIconDrawer from "./drawerComponents/StageIconDrawer";
+import StageGroupNode, {
+  StageGroupNodeName,
+} from "./flowComponents/StageGroupNode/StageGroupNode";
+import { StageNode, StageNodeName } from "./flowComponents/StageNode/StageNode";
+import { sqlElementsToLayout } from "./SqlLayoutService/SqlLayoutService";
 
 const options = { hideAttribution: true };
-const nodeTypes = { [StageNodeName]: StageNode };
+const nodeTypes = {
+  [StageNodeName]: StageNode,
+  [StageGroupNodeName]: StageGroupNode,
+};
+
+// this is a mock to allow for use of external configuration
+// const shouldUseGroupedLayout = true;
+const shouldUseGroupedLayout = false;
 
 const SqlFlow: FC<{ sparkSQL: EnrichedSparkSQL }> = ({
   sparkSQL,
@@ -31,35 +48,25 @@ const SqlFlow: FC<{ sparkSQL: EnrichedSparkSQL }> = ({
   const graphFilter = useAppSelector((state) => state.general.sqlMode);
   const selectedStage = useAppSelector((state) => state.general.selectedStage);
 
-  React.useEffect(() => {
-    if (!sparkSQL) return;
-    const { layoutNodes, layoutEdges } = SqlLayoutService.SqlElementsToLayout(
-      sparkSQL,
-      graphFilter,
-    );
-
-    setNodes(layoutNodes);
-  }, [sparkSQL.metricUpdateId]);
-
   useEffect(() => {
     if (!sparkSQL) return;
-    const { layoutNodes, layoutEdges } = SqlLayoutService.SqlElementsToLayout(
+
+    const { layoutNodes, layoutEdges } = sqlElementsToLayout(
       sparkSQL,
       graphFilter,
+      shouldUseGroupedLayout,
     );
 
     setNodes(layoutNodes);
     setEdges(layoutEdges);
-  }, [sparkSQL.uniqueId, graphFilter]);
+  }, [sparkSQL.metricUpdateId, sparkSQL.uniqueId, graphFilter]);
 
   useEffect(() => {
     if (instance) {
       setTimeout(instance.fitView, 20);
     }
   }, [instance, edges]);
 
-  useEffect(() => { }, [nodes]);
-
   const onConnect = useCallback(
     (params: any) =>
       setEdges((eds) =>
@@ -119,10 +126,11 @@ const SqlFlow: FC<{ sparkSQL: EnrichedSparkSQL }> = ({
           <Drawer
             anchor={"right"}
             open={selectedStage !== undefined}
-            onClose={() => dispatch(setSelectedStage({ selectedStage: undefined }))}
+            onClose={() =>
+              dispatch(setSelectedStage({ selectedStage: undefined }))
+            }
           >
-            <Box sx={{ minWidth: "400px" }}
-            >
+            <Box sx={{ minWidth: "400px" }}>
               <StageIconDrawer stage={selectedStage} />
             </Box>
           </Drawer>

spark-ui/src/components/SqlFlow/SqlLayoutService/SqlLayoutService.ts

@@ -0,0 +1,48 @@
+import { EnrichedSparkSQL, GraphFilter } from "../../../interfaces/AppStore";
+import { getLayoutElements } from "./dagreLayouts";
+import {
+  getFlowNodes,
+  getTopLevelNodes,
+  toFlowEdge,
+  transformEdgesToGroupEdges,
+} from "./layoutServiceBuilders";
+
+const getFlatFlowElements = (
+  sql: EnrichedSparkSQL,
+  graphFilter: GraphFilter,
+) => {
+  const { edges } = sql.filters[graphFilter];
+  const flowNodes = getFlowNodes(sql, graphFilter);
+  const flowEdges = edges.map(toFlowEdge);
+
+  return { flowNodes, flowEdges };
+};
+
+const getGroupedFlowElements = (
+  sql: EnrichedSparkSQL,
+  graphFilter: GraphFilter,
+) => {
+  const { edges } = sql.filters[graphFilter];
+
+  const flowNodes = getFlowNodes(sql, graphFilter);
+  const topLevelNodes = getTopLevelNodes(flowNodes);
+  const topLevelEdges = transformEdgesToGroupEdges(
+    flowNodes,
+    topLevelNodes,
+    edges,
+  );
+
+  return { flowNodes: topLevelNodes, flowEdges: topLevelEdges };
+};
+
+export function sqlElementsToLayout(
+  sql: EnrichedSparkSQL,
+  graphFilter: GraphFilter,
+  useGroupedLayout: boolean,
+) {
+  const { flowNodes, flowEdges } = useGroupedLayout
+    ? getGroupedFlowElements(sql, graphFilter)
+    : getFlatFlowElements(sql, graphFilter);
+
+  return getLayoutElements(flowNodes, flowEdges);
+}

spark-ui/src/components/SqlFlow/SqlLayoutService/dagreLayouts.ts

@@ -0,0 +1,52 @@
+import dagre from "dagre";
+import { Edge, Node, Position } from "reactflow";
+import { isNodeAGroup } from "../flowComponents/StageGroupNode/StageGroupNode";
+
+const buildDagreGraph = (rankdir: "LR" | "TB" = "LR") => {
+  const dagreGraph = new dagre.graphlib.Graph();
+  dagreGraph.setDefaultEdgeLabel(() => ({}));
+  dagreGraph.setGraph({ rankdir });
+
+  return dagreGraph;
+};
+
+const nodeSize = 280;
+export const nodeWidth = nodeSize;
+export const nodeHeight = nodeSize;
+const groupPadding = 40;
+const groupWidth = nodeWidth + groupPadding * 2;
+
+export const getLayoutElements = (
+  nodes: Node[],
+  edges: Edge[],
+): { layoutNodes: Node[]; layoutEdges: Edge[] } => {
+  const dagreGraph = buildDagreGraph();
+
+  nodes.forEach((node) => {
+    dagreGraph.setNode(node.id, {
+      width: isNodeAGroup(node) ? groupWidth : nodeWidth,
+      height: nodeHeight,
+    });
+  });
+
+  edges.forEach((edge) => {
+    dagreGraph.setEdge(edge.source, edge.target);
+  });
+
+  dagre.layout(dagreGraph);
+
+  nodes.forEach((node) => {
+    const nodeWithPosition = dagreGraph.node(node.id);
+    node.targetPosition = Position.Left;
+    node.sourcePosition = Position.Right;
+
+    // We are shifting the dagre node position (anchor=center center) to the top left
+    // so it matches the React Flow node anchor point (top left).
+    node.position = {
+      x: nodeWithPosition.x - nodeWidth / 2,
+      y: nodeWithPosition.y - nodeHeight / 2,
+    };
+  });
+
+  return { layoutNodes: nodes, layoutEdges: edges };
+};