feat(ui): add option to show visualized query stats

pinot-controller/src/main/resources/app/components/Query/VisualizeQueryStageStats.tsx

@@ -0,0 +1,201 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React from "react";
+import ReactFlow, { Background, Controls, MiniMap, Handle, Node, Edge } from "react-flow-renderer";
+import dagre from "dagre";
+import { Typography, useTheme } from "@material-ui/core";
+import "react-flow-renderer/dist/style.css";
+import isEmpty from "lodash/isEmpty";
+
+/**
+ * Main component to visualize query stage stats as a flowchart.
+ */
+export const VisualizeQueryStageStats = ({ stageStats }) => {
+  const { nodes, edges } = generateFlowElements(stageStats); // Generate nodes and edges from input data
+
+  if(isEmpty(stageStats)) {
+    return (
+      <Typography style={{height: "100px", textAlign: "center"}} variant="body1">
+        No stats available
+      </Typography>
+    );
+  }
+
+  return (
+    <div style={{ height: 500 }}>
+      <ReactFlow
+        nodes={nodes}
+        edges={edges}
+        fitView
+        nodeTypes={nodeTypes} // Use custom node types
+        zoomOnScroll={false}
+      >
+        <Background />
+        <Controls showInteractive={false} />
+        <MiniMap />
+      </ReactFlow>
+    </div>
+  );
+};
+
+// ------------------------------------------------------------
+// Helper functions and constants
+
+// Constants for node styling and layout
+const NODE_PADDING = 10;
+const NODE_HEADER_CONTENT_MARGIN = 5;
+const NODE_HEADER_HEIGHT = 20;
+const NODE_CONTENT_HEIGHT = 16;
+const NODE_TEXT_CHAR_WIDTH = 7;
+const NODE_MIN_WIDTH = 150;
+
+/**
+ * Calculates the dimensions of a node based on its content.
+ */
+const calculateNodeDimensions = (data) => {
+  const contentWidth = Math.max(
+    Object.entries(data)
+      .map(([key, value]) => key.length + String(value).length + 1) // Estimate character count in the content
+      .reduce((max, length) => Math.max(max, length), 0) * NODE_TEXT_CHAR_WIDTH,
+    NODE_MIN_WIDTH
+  );
+
+  const contentHeight =
+    Object.keys(data).length * NODE_CONTENT_HEIGHT + // Height for each data line
+    NODE_HEADER_HEIGHT + // Height for the header
+    NODE_HEADER_CONTENT_MARGIN * 2 + // Margin between header and content
+    NODE_PADDING * 2; // Padding around the node
+
+  return { width: contentWidth + NODE_PADDING * 2, height: contentHeight };
+};
+
+/**
+ * Applies Dagre layout to position nodes and edges.
+ */
+const layoutNodesAndEdges = (nodes, edges, direction = "TB") => {
+  const dagreGraph = new dagre.graphlib.Graph();
+  dagreGraph.setDefaultEdgeLabel(() => ({})); // Default edge properties
+  dagreGraph.setGraph({ rankdir: direction }); // Layout direction
+
+  // Add nodes to the graph
+  nodes.forEach((node) => {
+    dagreGraph.setNode(node.id, { width: node.width, height: node.height });
+  });
+
+  // Add edges to the graph
+  edges.forEach((edge) => {
+    dagreGraph.setEdge(edge.source, edge.target);
+  });
+
+  // Perform Dagre layout
+  dagre.layout(dagreGraph);
+
+  const isHorizontal = direction === "LR";
+  return {
+    nodes: nodes.map((node) => {
+      const layoutedNode = dagreGraph.node(node.id); // Get node's position
+      return {
+        ...node,
+        position: {
+          x: layoutedNode.x - node.width / 2, // Center node horizontally
+          y: layoutedNode.y - node.height / 2, // Center node vertically
+        },
+        targetPosition: isHorizontal ? "left" : "top",
+        sourcePosition: isHorizontal ? "right" : "bottom",
+      };
+    }),
+    edges,
+  };
+};
+
+/**
+ * Recursively generates nodes and edges for the flowchart from a hierarchical data structure.
+ */
+const generateFlowElements = (stats) => {
+  const nodes: Node[] = [];
+  const edges: Edge[] = [];
+
+  /**
+   * Traverses the hierarchy and builds nodes and edges.
+   */
+  const traverseTree = (node, level, index, parentId) => {
+    const { children, ...data } = node;
+    const id = `${level}-${index}`; // Unique ID for the node
+    const { width, height } = calculateNodeDimensions(data);
+
+    // Add the node
+    nodes.push({ id, type: "customNode", data, position: { x: 0, y: 0 }, width, height });
+
+    // Add an edge if this node has a parent
+    if (parentId) {
+      edges.push({ id: `edge-${parentId}-${id}`, source: parentId, target: id });
+    }
+
+    // Recursively process children
+    children?.forEach((child, idx) => traverseTree(child, level + 1, index + idx, id));
+  };
+
+  traverseTree(stats, 0, 0, null); // Start traversal from the root node
+  return layoutNodesAndEdges(nodes, edges);
+};
+
+/**
+ * Custom Node Renderer for React Flow.
+ */
+const CustomNode = ({ data, ...props }) => {
+  const theme = useTheme();
+  return (
+    <div
+      style={{
+        border: "1px solid #ccc",
+        borderRadius: "8px",
+        backgroundColor: "#fff",
+        padding: NODE_PADDING,
+        boxShadow: "0 2px 5px rgba(0, 0, 0, 0.1)",
+        minWidth: NODE_MIN_WIDTH,
+      }}
+    >
+      <Handle type="source" position={props.sourcePosition} />
+      <div
+        style={{
+          fontWeight: "bold",
+          color: theme.palette.primary.main,
+          fontSize: "18px",
+          lineHeight: `${NODE_HEADER_HEIGHT}px`,
+        }}
+      >
+        {data.type || "Unknown Type"} {/* Display node type */}
+      </div>
+      <hr style={{ margin: NODE_HEADER_CONTENT_MARGIN, color: "#ddd" }} />
+      <div style={{ fontSize: "14px", lineHeight: `${NODE_CONTENT_HEIGHT}px` }}>
+        {Object.entries(data).map(([key, value]) => (
+          <div key={key}>
+            <strong>{key}:</strong> {String(value)} {/* Display key-value pairs */}
+          </div>
+        ))}
+      </div>
+      <Handle type="target" position={props.targetPosition} />
+    </div>
+  );
+};
+
+// Define custom node types for React Flow
+const nodeTypes = {
+  customNode: CustomNode,
+};

pinot-controller/src/main/resources/app/pages/Query.tsx

@@ -20,7 +20,7 @@
 
 import React, { useEffect, useState } from 'react';
 import { makeStyles } from '@material-ui/core/styles';
-import { Grid, Checkbox, Button, FormControl, Input, InputLabel, Box, Typography } from '@material-ui/core';
+import { Grid, Checkbox, Button, FormControl, Input, InputLabel, Box, Typography, ButtonGroup } from '@material-ui/core';
 import Alert from '@material-ui/lab/Alert';
 import FileCopyIcon from '@material-ui/icons/FileCopy';
 import { SqlException, TableData } from 'Models';
@@ -48,6 +48,13 @@ import '../styles/styles.css';
 import {Resizable} from "re-resizable";
 import { useHistory, useLocation } from 'react-router';
 import sqlFormatter from '@sqltools/formatter';
+import { VisualizeQueryStageStats } from '../components/Query/VisualizeQueryStageStats';
+
+enum ResultViewType {
+  TABULAR = 'tabular',
+  JSON = 'json',
+  VISUAL = 'visual',
+}
 
 const useStyles = makeStyles((theme) => ({
   title: {
@@ -200,13 +207,14 @@ const QueryPage = () => {
     columns: [],
     records: [],
   });
+  const [resultViewType, setResultViewType] = useState(ResultViewType.TABULAR);
+  const [stageStats, setStageStats] = useState({});
 
   const [warnings, setWarnings] = useState<Array<string>>([]);
 
   const [checked, setChecked] = React.useState({
     tracing: queryParam.get('tracing') === 'true',
     useMSE: queryParam.get('useMSE') === 'true',
-    showResultJSON: false,
   });
 
   const queryExecuted = React.useRef(false);
@@ -325,6 +333,7 @@ const QueryPage = () => {
     setResultData(results.result || { columns: [], records: [] });
     setQueryStats(results.queryStats || { columns: responseStatCols, records: [] });
     setOutputResult(JSON.stringify(results.data, null, 2) || '');
+    setStageStats(results?.data?.stageStats || {});
     setWarnings(extractWarnings(results));
     setQueryLoader(false);
     queryExecuted.current = false;
@@ -405,8 +414,7 @@ const QueryPage = () => {
       setInputQuery(query);
       setChecked({
         tracing: queryParam.get('tracing') === 'true',
-        useMSE: queryParam.get('useMse') === 'true',
-        showResultJSON: checked.showResultJSON,
+        useMSE: queryParam.get('useMse') === 'true'
       });
       setQueryTimeout(Number(queryParam.get('timeout') || '') || '');
       setBoolFlag(!boolFlag);
@@ -663,27 +671,29 @@ const QueryPage = () => {
                         ) : null}
 
                         <FormControlLabel
+                          labelPlacement='start'
                           control={
-                            <Switch
-                              checked={checked.showResultJSON}
-                              onChange={handleChange}
-                              name="showResultJSON"
-                              color="primary"
-                            />
+                            <ButtonGroup color='primary' size='small'>
+                              <Button onClick={() => setResultViewType(ResultViewType.TABULAR)} variant={resultViewType === ResultViewType.TABULAR ? "contained" : "outlined"}>Tabular</Button>
+                              <Button onClick={() => setResultViewType(ResultViewType.JSON)} variant={resultViewType === ResultViewType.JSON ? "contained" : "outlined"}>Json</Button>
+                              <Button onClick={() => setResultViewType(ResultViewType.VISUAL)} variant={resultViewType === ResultViewType.VISUAL ? "contained" : "outlined"}>Visual</Button>
+                            </ButtonGroup>
                           }
-                          label="Show JSON format"
+                          label={<Typography style={{marginRight: "8px"}}>View</Typography>}
+                          style={{marginRight: 0}}
                           className={classes.runNowBtn}
                         />
                       </Grid>
-                      {!checked.showResultJSON ? (
+                      {resultViewType === ResultViewType.TABULAR && (
                         <CustomizedTables
                           title="Query Result"
                           data={resultData}
                           isSticky={true}
                           showSearchBox={true}
                           inAccordionFormat={true}
                         />
-                      ) : resultData.columns.length ? (
+                      )} 
+                      {resultViewType === ResultViewType.JSON && (
                         <SimpleAccordion
                           headerTitle="Query Result (JSON Format)"
                           showSearchBox={false}
@@ -695,7 +705,15 @@ const QueryPage = () => {
                             autoCursor={false}
                           />
                         </SimpleAccordion>
-                      ) : null}
+                      )}
+                      {resultViewType === ResultViewType.VISUAL && (
+                        <SimpleAccordion
+                          headerTitle="Query Stats Visualized"
+                          showSearchBox={false}
+                        >
+                          <VisualizeQueryStageStats stageStats={stageStats} />
+                        </SimpleAccordion>
+                      )}
                     </>
                   ) : null}
                 </Grid>