GraphVisualizer

gistfile1.java

package org.apache.drill.common.graph;

import org.apache.drill.common.logical.LogicalPlan;
import org.apache.drill.common.logical.data.LogicalOperator;
import org.apache.drill.common.logical.data.Scan;
import org.w3c.dom.Document;

import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.awt.*;
import java.io.FileOutputStream;
import java.util.*;
import com.mxgraph.layout.hierarchical.mxHierarchicalLayout;
import com.mxgraph.util.mxCellRenderer;
import com.mxgraph.util.mxConstants;
import com.mxgraph.util.mxUtils;
import com.mxgraph.view.mxGraph;

public class GraphVisualizer {
  
  private static final int VERTEX_HEIGHT = 30;
  private static final int VERTEX_WIDTH = 140;
  private static final int SCAN_WIDTH = 200;
  private static final int VERTEX_FONTSIZE = 6;
  private static final String STYLE_SCAN = "scan";
  
  private static mxGraph buildDirectedGraphMX(LogicalPlan plan) {
    Collection<Edge<AdjacencyList<LogicalOperator>.Node>> edges = plan.getGraph().getAdjList().getAllEdges();    
    mxGraph mx = new mxGraph();
    mx.getModel().beginUpdate();
    initStyle(mx);
    Object parent = mx.getDefaultParent();
    Map<LogicalOperator, Object> vertexes = new HashMap<>();    
    for(Edge<AdjacencyList<LogicalOperator>.Node> edge:edges){
      LogicalOperator node = edge.from.getNodeValue();
      if(!vertexes.containsKey(node)){
        Object cell = insertVertex(parent, node, mx);
        vertexes.put(node, cell);
      }
      Object source = vertexes.get(node);
      node = edge.to.getNodeValue();
      if(!vertexes.containsKey(node)){
        
        Object cell = insertVertex(parent, node, mx);      
        vertexes.put(node, cell);
      }
      Object target = vertexes.get(node);
      mx.insertEdge(parent, null, null, source, target);
    }
    mx.getModel().endUpdate();
    vertexes.clear();
    return mx;
  }

  private static void initStyle(mxGraph mx) {
    //style
    Map<String, Object> scanStyle = new HashMap<>();
    String color = "#"+mxUtils.getHexColorString(new Color(255, 195, 217, 0));
    scanStyle.put(mxConstants.STYLE_FILLCOLOR, color);//
    mx.getStylesheet().putCellStyle(STYLE_SCAN, scanStyle);
  }

  private static Object insertVertex(Object parent, LogicalOperator operator, mxGraph mx) {
    String style = null;
    if(operator instanceof Scan){
      style = STYLE_SCAN;
//      style = null;
    }
    return mx.insertVertex(parent, null, operator, 0, 0, VERTEX_WIDTH, VERTEX_HEIGHT, style);
  }


  static int WIDTH = 1200;
  static int HEIGHT = 1200;
  
  public static void visualizeMX(LogicalPlan plan, String svgPath){
    try {
      mxGraph graph = buildDirectedGraphMX(plan);
      mxHierarchicalLayout hierarchical = new mxHierarchicalLayout(graph);
//      mxCompactTreeLayout tree =  new mxCompactTreeLayout(graph, true);
//      tree.setEdgeRouting(true);
//      tree.execute(graph.getDefaultParent());
      hierarchical.setInterRankCellSpacing(100.0);
      hierarchical.execute(graph.getDefaultParent());
      Document svg = mxCellRenderer.createSvgDocument(graph, null, 1, Color.WHITE, null);
      TransformerFactory tFactory = TransformerFactory.newInstance();
      Transformer transformer = tFactory.newTransformer();
      
      DOMSource source = new DOMSource(svg);
      StreamResult result = new StreamResult(new FileOutputStream(svgPath));
      transformer.transform(source, result); 
      
    }catch(Exception e){
      e.printStackTrace();  //e:
    }
  }
  
}