Implement graphviz printer for SampleNode

trinodb · Oct 24, 2014 · af0879e · af0879e
1 parent 617ca77
commit af0879e
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java b/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java
@@ -31,6 +31,7 @@
 import com.facebook.presto.sql.planner.plan.PlanVisitor;
 import com.facebook.presto.sql.planner.plan.ProjectNode;
 import com.facebook.presto.sql.planner.plan.RowNumberNode;
+import com.facebook.presto.sql.planner.plan.SampleNode;
 import com.facebook.presto.sql.planner.plan.SemiJoinNode;
 import com.facebook.presto.sql.planner.plan.SinkNode;
 import com.facebook.presto.sql.planner.plan.SortNode;
@@ -79,6 +80,7 @@ private enum NodeType
         WINDOW,
         UNION,
         SORT,
+        SAMPLE,
         MARK_DISTINCT,
         MATERIALIZE_SAMPLE,
         INDEX_SOURCE,
@@ -104,6 +106,7 @@ private enum NodeType
             .put(NodeType.MATERIALIZE_SAMPLE, "hotpink")
             .put(NodeType.INDEX_SOURCE, "dodgerblue3")
             .put(NodeType.UNNEST, "crimson")
+            .put(NodeType.SAMPLE, "goldenrod4")
             .build());
 
     static {
@@ -204,6 +207,13 @@ protected Void visitPlan(PlanNode node, Void context)
             throw new UnsupportedOperationException(format("Node %s does not have a Graphviz visitor", node.getClass().getName()));
         }
 
+        @Override
+        public Void visitSample(SampleNode node, Void context)
+        {
+            printNode(node, format("Sample[type=%s, ratio=%f, rescaled=%s]", node.getSampleType(), node.getSampleRatio(), node.isRescaled()), NODE_COLORS.get(NodeType.SAMPLE));
+            return node.getSource().accept(this, context);
+        }
+
         @Override
         public Void visitSort(SortNode node, Void context)
         {