diff options
author | Jinfeng Ni <jni@maprtech.com> | 2015-01-20 09:07:33 -0800 |
---|---|---|
committer | Jinfeng Ni <jni@maprtech.com> | 2015-01-20 12:08:35 -0800 |
commit | a8036d29e0d50e4753b152803b13e1337800c67c (patch) | |
tree | cba9bb4fe3986c4593955c09b5f93990588a9692 /exec/java-exec | |
parent | 8103099e02fea0a93788155c1f7e95532b9d82f2 (diff) |
DRILL-2030: fix column prefix for CTAS when query has select * and expression.
Code refactor.
Diffstat (limited to 'exec/java-exec')
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() |