DRILL-2030: fix column prefix for CTAS when query has select * and expression.

Code refactor.

3 files changed, 49 insertions, 17 deletions

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/visitor/StarColumnConverter.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/visitor/StarColumnConverter.java
index 1511e4936..fcaafc8b9 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/visitor/StarColumnConverter.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/physical/visitor/StarColumnConverter.java
@@ -29,6 +29,8 @@ import org.apache.drill.exec.planner.physical.Prel;
 import org.apache.drill.exec.planner.physical.ProjectAllowDupPrel;
 import org.apache.drill.exec.planner.physical.ProjectPrel;
 import org.apache.drill.exec.planner.physical.ScanPrel;
+import org.apache.drill.exec.planner.physical.ScreenPrel;
+import org.apache.drill.exec.planner.physical.WriterPrel;
 import org.eigenbase.rel.RelNode;
 import org.eigenbase.reltype.RelDataType;
 import org.eigenbase.reltype.RelDataTypeField;
@@ -45,7 +47,7 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
 
   private static final AtomicLong tableNumber = new AtomicLong(0);
 
-  public static Prel insertRenameProject(Prel root, RelDataType origRowType) {
+  public static Prel insertRenameProject(Prel root) {
     // Prefixing columns for columns expanded from star column :
     // Insert one project under screen (PUS) to remove prefix, and one project above scan (PAS) to add prefix.
     // PUS AND PAS are required, when
@@ -58,10 +60,29 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
     boolean [] prefixedForStar = new boolean [1];
     prefixedForStar[0] = false;
 
-    //root should be screen / writer : no need to rename for the root.
-    Prel child = ((Prel) root.getInput(0)).accept(INSTANCE, prefixedForStar);
+    return root.accept(INSTANCE, prefixedForStar);
+  }
+
+  @Override
+  public Prel visitScreen(ScreenPrel prel, boolean[] prefixedForStar) throws RuntimeException {
+    return insertProjUnderScreen(prel, prefixedForStar, prel.getChild().getRowType());
+  }
 
-    int fieldCount = root.getRowType().isStruct()? root.getRowType().getFieldCount():1;
+  @Override
+  public Prel visitWriter(WriterPrel prel, boolean[] prefixedForStar) throws RuntimeException {
+    Prel newPrel = insertProjUnderScreen(prel, prefixedForStar, prel.getChild().getRowType());
+
+    prefixedForStar[0] = false;
+
+    return newPrel;
+  }
+
+  // insert PUS: Project Under Screen, when necessary.
+  private Prel insertProjUnderScreen(Prel prel, boolean[] prefixedForStar, RelDataType origRowType) {
+
+    Prel child = ((Prel) prel.getInput(0)).accept(INSTANCE, prefixedForStar);
+
+    ProjectPrel proj = null;
 
     if (prefixedForStar[0]) {
       List<RexNode> exprs = Lists.newArrayList();
@@ -72,27 +93,31 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
 
       RelDataType newRowType = RexUtil.createStructType(child.getCluster().getTypeFactory(), exprs, origRowType.getFieldNames());
 
-      // Insert PUS : remove the prefix and keep the orignal field name.
+      int fieldCount = prel.getRowType().isStruct()? prel.getRowType().getFieldCount():1;
+
+      // Insert PUS : remove the prefix and keep the original field name.
       if (fieldCount > 1) { // // no point in allowing duplicates if we only have one column
-        child = new ProjectAllowDupPrel(child.getCluster(), child.getTraitSet(), child, exprs, newRowType);
+        proj = new ProjectAllowDupPrel(child.getCluster(), child.getTraitSet(), child, exprs, newRowType);
       } else {
-        child = new ProjectPrel(child.getCluster(), child.getTraitSet(), child, exprs, newRowType);
+        proj = new ProjectPrel(child.getCluster(), child.getTraitSet(), child, exprs, newRowType);
       }
 
       List<RelNode> children = Lists.newArrayList();
-      children.add( child);
-      return (Prel) root.copy(root.getTraitSet(), children);
 
+      children.add(proj);
+      return (Prel) prel.copy(prel.getTraitSet(), children);
     } else {
-      return root;
+      return prel;
     }
 
   }
 
   @Override
   public Prel visitPrel(Prel prel, boolean [] prefixedForStar) throws RuntimeException {
-    // there exists other expression, in addition to a star column: require prefixRename.
-    if (StarColumnHelper.containsStarColumn(prel.getRowType()) && prel.getRowType().getFieldNames().size() > 1) {
+    // Require prefix rename : there exists other expression, in addition to a star column.
+    if (!prefixedForStar[0]  // not set yet.
+        && StarColumnHelper.containsStarColumn(prel.getRowType())
+        && prel.getRowType().getFieldNames().size() > 1) {
       prefixedForStar[0] = true;
     }
 
@@ -103,8 +128,8 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
     }
 
     // For project, we need make sure that the project's field name is same as the input,
-    // when the project expression is RexInPutRef. This is necessary since Optiq may use
-    // an arbitrary name for the project's field name.
+    // when the project expression is RexInPutRef, since we may insert a PAS which will
+    // rename the projected fields.
     if (prel instanceof ProjectPrel) {
 
       ProjectPrel proj = (ProjectPrel) prel;
@@ -122,7 +147,7 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
         }
       }
 
-      // Make sure the field names are unique : Optiq does not allow duplicate field names in a rowType.
+      // Make sure the field names are unique : no allow of duplicate field names in a rowType.
       fieldNames = makeUniqueNames(fieldNames);
 
       RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(), proj.getProjects(), fieldNames);
@@ -166,7 +191,6 @@ public class StarColumnConverter extends BasePrelVisitor<Prel, boolean[], Runtim
     }
   }
 
-
   private List<String> makeUniqueNames(List<String> names) {
 
     // We have to search the set of original names, plus the set of unique names that will be used finally .
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
index 83d2d0bf6..79603ebd8 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/handlers/DefaultSqlHandler.java
@@ -187,7 +187,7 @@ public class DefaultSqlHandler extends AbstractSqlHandler {
      * duplicate columns, since user could "explicitly" ask for duplicate columns ( select *, col, *).
      * The rest of projects will remove the duplicate column when we generate POP in json format.
      */
-    phyRelNode = StarColumnConverter.insertRenameProject(phyRelNode, phyRelNode.getRowType());
+    phyRelNode = StarColumnConverter.insertRenameProject(phyRelNode);
 
     /*
      * 1.)
diff --git a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/writer/TestParquetWriter.java b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/writer/TestParquetWriter.java
index 8ecb6c17a..dc80afedf 100644
--- a/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/writer/TestParquetWriter.java
+++ b/exec/java-exec/src/test/java/org/apache/drill/exec/physical/impl/writer/TestParquetWriter.java
@@ -224,6 +224,14 @@ public class TestParquetWriter extends BaseTestQuery {
     runTestAndValidate(selection, validateSelection, inputTable, "region_boolean_parquet");
   }
 
+  @Test //DRILL-2030
+  public void testWriterWithStarAndExp() throws Exception {
+    String selection = " *, r_regionkey + 1";
+    String validateSelection = "r_regionkey, r_name, r_comment, r_regionkey + 1";
+    String inputTable = "cp.`tpch/region.parquet`";
+    runTestAndValidate(selection, validateSelection, inputTable, "region_star_exp");
+  }
+
   public void compareParquetReadersColumnar(String selection, String table) throws Exception {
     String query = "select " + selection + " from " + table;
     testBuilder()