DRILL-6476: Generate explain plan which shows relation between Latera…

…l and the corresponding Unnest.

closes #1308

exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/PhysicalPlanCreator.java

@@ -52,11 +52,6 @@ public QueryContext getContext() {
     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()));

exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/UnnestPrel.java

@@ -75,4 +75,7 @@ public boolean needsFinalColumnReordering() {
     return true;
   }
 
+  public Class<?> getParentClass() {
+    return CorrelatePrel.class;
+  }
 }

exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/explain/NumberingRelWriter.java

@@ -19,6 +19,7 @@
 
 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.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,23 +59,18 @@ public NumberingRelWriter(Map<Prel, OpId> ids, PrintWriter pw, SqlExplainLevel d
     this.pw = pw;
     this.ids = ids;
     this.detailLevel = detailLevel;
+    this.sourceOperatorRegistry = new HashMap<>();
   }
 
   //~ Methods ----------------------------------------------------------------
 
   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 @@ protected void explain_(
     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 @@ protected void explain_(
     }
     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) {

exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/lateraljoin/TestLateralPlans.java

@@ -435,4 +435,25 @@ private String getRightChildOfLateral(String explain) throws Exception {
     String CorrelateUnnest = matcher.group(0);
     return CorrelateUnnest.substring(CorrelateUnnest.lastIndexOf("Scan"));
   }
+
+
+  //The following test is for testing the explain plan contains relation between lateral and corresponding unnest.
+  @Test
+  public void testLateralAndUnnestExplainPlan() throws Exception {
+    String Sql = "select c.* from cp.`lateraljoin/nested-customer.json` c, unnest(c.orders) Orders(ord)";
+    ClusterFixtureBuilder builder = ClusterFixture.builder(dirTestWatcher)
+            .setOptionDefault(ExecConstants.ENABLE_UNNEST_LATERAL_KEY, true)
+            .setOptionDefault(ExecConstants.SLICE_TARGET, 1);
+
+    try (ClusterFixture cluster = builder.build();
+         ClientFixture client = cluster.clientFixture()) {
+      String explain = client.queryBuilder().sql(Sql).explainText();
+      String srcOp = explain.substring(explain.indexOf("srcOp"));
+      assertTrue(srcOp != null && srcOp.length() > 0);
+      String correlateFragmentPattern = srcOp.substring(srcOp.indexOf("=")+1, srcOp.indexOf("]"));
+      assertTrue(correlateFragmentPattern != null && correlateFragmentPattern.length() > 0);
+      Matcher matcher = Pattern.compile(correlateFragmentPattern + ".*Correlate", Pattern.MULTILINE | Pattern.DOTALL).matcher(explain);
+      assertTrue(matcher.find());
+    }
+  }
 }