DRILL-6476: Generate explain plan which shows relation between Lateral and the corresponding Unnest.

closes #1308

3 files changed, 61 insertions, 16 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/PhysicalPlanCreator.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/PhysicalPlanCreator.java
index 6a94662e7..220add66f 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/PhysicalPlanCreator.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/PhysicalPlanCreator.java
@@ -52,11 +52,6 @@ public class PhysicalPlanCreator {
     return context;
   }
 
-//  public int getOperatorId(Prel prel){
-//    OpId id = opIdMap.get(prel);
-//    return id.getAsSingleInt();
-//  }
-
   public PhysicalOperator addMetadata(Prel originalPrel, PhysicalOperator op){
     op.setOperatorId(opIdMap.get(originalPrel).getAsSingleInt());
     op.setCost(originalPrel.estimateRowCount(originalPrel.getCluster().getMetadataQuery()));
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/UnnestPrel.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/UnnestPrel.java
index cd598eb98..a22beea0e 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/UnnestPrel.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/UnnestPrel.java
@@ -75,4 +75,7 @@ public class UnnestPrel extends DrillUnnestRelBase implements Prel {
     return true;
   }
 
+  public Class<?> getParentClass() {
+    return CorrelatePrel.class;
+  }
 }
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/explain/NumberingRelWriter.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/explain/NumberingRelWriter.java
index 045dba9d4..38b97b67d 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/explain/NumberingRelWriter.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/explain/NumberingRelWriter.java
@@ -19,6 +19,7 @@ package org.apache.drill.exec.planner.physical.explain;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -31,8 +32,10 @@ import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.runtime.FlatLists;
 import org.apache.calcite.sql.SqlExplainLevel;
 import org.apache.calcite.util.Pair;
+import org.apache.drill.exec.planner.physical.CorrelatePrel;
 import org.apache.drill.exec.planner.physical.HashJoinPrel;
 import org.apache.drill.exec.planner.physical.Prel;
+import org.apache.drill.exec.planner.physical.UnnestPrel;
 import org.apache.drill.exec.planner.physical.explain.PrelSequencer.OpId;
 
 import com.google.common.collect.ImmutableList;
@@ -47,6 +50,7 @@ class NumberingRelWriter implements RelWriter {
   private final SqlExplainLevel detailLevel;
   protected final Spacer spacer = new Spacer();
   private final List<Pair<String, Object>> values = new ArrayList<>();
+  private final Map<String, Prel> sourceOperatorRegistry;
 
   private final Map<Prel, OpId> ids;
   //~ Constructors -----------------------------------------------------------
@@ -55,6 +59,7 @@ class NumberingRelWriter implements RelWriter {
     this.pw = pw;
     this.ids = ids;
     this.detailLevel = detailLevel;
+    this.sourceOperatorRegistry = new HashMap<>();
   }
 
   //~ Methods ----------------------------------------------------------------
@@ -62,16 +67,10 @@ class NumberingRelWriter implements RelWriter {
   protected void explain_(
       RelNode rel,
       List<Pair<String, Object>> values) {
-    List<RelNode> inputs = rel.getInputs();
-    if (rel instanceof HashJoinPrel && ((HashJoinPrel) rel).isSwapped()) {
-      HashJoinPrel joinPrel = (HashJoinPrel) rel;
-      inputs = FlatLists.of(joinPrel.getRight(), joinPrel.getLeft());
-    }
-
     RelMetadataQuery mq = RelMetadataQuery.instance();
     if (!mq.isVisibleInExplain(rel, detailLevel)) {
       // render children in place of this, at same level
-      explainInputs(inputs);
+      explainInputs(rel);
       return;
     }
 
@@ -95,6 +94,7 @@ class NumberingRelWriter implements RelWriter {
     s.append(rel.getRelTypeName().replace("Prel", ""));
     if (detailLevel != SqlExplainLevel.NO_ATTRIBUTES) {
       int j = 0;
+      s.append(getDependentSrcOp(rel));
       for (Pair<String, Object> value : values) {
         if (value.right instanceof RelNode) {
           continue;
@@ -125,14 +125,61 @@ class NumberingRelWriter implements RelWriter {
     }
     pw.println(s);
     spacer.add(2);
-    explainInputs(inputs);
+    explainInputs(rel);
     spacer.subtract(2);
   }
 
-  private void explainInputs(List<RelNode> inputs) {
-    for (RelNode input : inputs) {
-      input.explain(this);
+  private String getDependentSrcOp(RelNode rel) {
+    if (rel instanceof UnnestPrel) {
+      return this.getDependentSrcOp((UnnestPrel) rel);
+    }
+    return "";
+  }
+
+  private String getDependentSrcOp(UnnestPrel unnest) {
+    Prel parent = this.getRegisteredPrel(unnest.getParentClass());
+    if (parent != null && parent instanceof CorrelatePrel) {
+      OpId id = ids.get(parent);
+      return String.format(" [srcOp=%02d-%02d] ", id.fragmentId, id.opId);
     }
+    return "";
+  }
+
+  public void register(Prel toRegister) {
+    this.sourceOperatorRegistry.put(toRegister.getClass().getSimpleName(), toRegister);
+  }
+
+  public Prel getRegisteredPrel(Class<?> classname) {
+    return this.sourceOperatorRegistry.get(classname.getSimpleName());
+  }
+
+  public void unRegister(Prel unregister) {
+    this.sourceOperatorRegistry.remove(unregister.getClass().getSimpleName());
+  }
+
+
+  private void explainInputs(RelNode rel) {
+    if (rel instanceof CorrelatePrel) {
+      this.explainInputs((CorrelatePrel) rel);
+    } else {
+      List<RelNode> inputs = rel.getInputs();
+      if (rel instanceof HashJoinPrel && ((HashJoinPrel) rel).isSwapped()) {
+        HashJoinPrel joinPrel = (HashJoinPrel) rel;
+        inputs = FlatLists.of(joinPrel.getRight(), joinPrel.getLeft());
+      }
+      for (RelNode input : inputs) {
+        input.explain(this);
+      }
+    }
+  }
+
+  //Correlate is handled differently because explain plan
+  //needs to show relation between Lateral and Unnest operators.
+  private void explainInputs(CorrelatePrel correlate) {
+    correlate.getInput(0).explain(this);
+    this.register(correlate);
+    correlate.getInput(1).explain(this);
+    this.unRegister(correlate);
   }
 
   public final void explain(RelNode rel, List<Pair<String, Object>> valueList) {