[calcite] branch master updated: [CALCITE-3239] Calc#accept(RexShuttle shuttle) does not update rowType. (Jin Xing)

[hyuan]

This is an automated email from the ASF dual-hosted git repository.

hyuan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git


The following commit(s) were added to refs/heads/master by this push:
     new be01394  [CALCITE-3239] Calc#accept(RexShuttle shuttle) does not 
update rowType. (Jin Xing)
be01394 is described below

commit be01394210102ea83dc83f95ce5a59c60525b8d3
Author: jinxing <jinxing.co...@gmail.com>
AuthorDate: Mon Aug 12 00:04:31 2019 +0800

    [CALCITE-3239] Calc#accept(RexShuttle shuttle) does not update rowType. 
(Jin Xing)
---
 .../java/org/apache/calcite/rel/core/Calc.java     | 22 ++++++--
 .../org/apache/calcite/test/RexShuttleTest.java    | 60 ++++++++++++++++++++++
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/rel/core/Calc.java 
b/core/src/main/java/org/apache/calcite/rel/core/Calc.java
index 278fba8..27cd0df 100644
--- a/core/src/main/java/org/apache/calcite/rel/core/Calc.java
+++ b/core/src/main/java/org/apache/calcite/rel/core/Calc.java
@@ -27,10 +27,14 @@ import org.apache.calcite.rel.RelWriter;
 import org.apache.calcite.rel.SingleRel;
 import org.apache.calcite.rel.metadata.RelMdUtil;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexLocalRef;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexProgram;
+import org.apache.calcite.rex.RexProgramBuilder;
 import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.util.Litmus;
 import org.apache.calcite.util.Util;
 
@@ -167,12 +171,20 @@ public abstract class Calc extends SingleRel {
         && condition == oldCondition) {
       return this;
     }
-    return copy(traitSet, getInput(),
-        new RexProgram(program.getInputRowType(),
-            exprs,
+
+    final RexBuilder rexBuilder = getCluster().getRexBuilder();
+    final RelDataType rowType =
+        RexUtil.createStructType(
+            rexBuilder.getTypeFactory(),
             projects,
-            (RexLocalRef) condition,
-            program.getOutputRowType()));
+            this.rowType.getFieldNames(),
+            null);
+    final RexProgram newProgram =
+        RexProgramBuilder.create(
+            rexBuilder, program.getInputRowType(), exprs, projects,
+            condition, rowType, true, null)
+        .getProgram(false);
+    return copy(traitSet, getInput(), newProgram);
   }
 }
 
diff --git a/core/src/test/java/org/apache/calcite/test/RexShuttleTest.java 
b/core/src/test/java/org/apache/calcite/test/RexShuttleTest.java
index da90d60..f3edfb6 100644
--- a/core/src/test/java/org/apache/calcite/test/RexShuttleTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RexShuttleTest.java
@@ -16,11 +16,19 @@
  */
 package org.apache.calcite.test;
 
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgram;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.logical.LogicalCalc;
+import org.apache.calcite.rel.rules.ProjectToCalcRule;
 import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLocalRef;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.tools.RelBuilder;
 
@@ -71,6 +79,58 @@ public class RexShuttleTest {
 
     assertThat(type, is(type2));
   }
+
+  @Test
+  public void testCalcUpdatesRowType() {
+    final RelBuilder builder = 
RelBuilder.create(RelBuilderTest.config().build());
+
+    // Equivalent SQL: SELECT deptno, sal, sal + 20 FROM emp
+    final RelNode root =
+        builder
+            .scan("EMP")
+            .project(
+                builder.field("DEPTNO"),
+                builder.field("SAL"),
+                builder.call(SqlStdOperatorTable.PLUS,
+                    builder.field("SAL"), builder.literal(20)))
+            .build();
+
+    HepProgram program = new HepProgramBuilder()
+        .addRuleInstance(ProjectToCalcRule.INSTANCE)
+        .build();
+    HepPlanner planner = new HepPlanner(program);
+    planner.setRoot(root);
+    LogicalCalc calc = (LogicalCalc) planner.findBestExp();
+
+    final RelNode calcWithCastViaRexShuttle = calc.accept(new RexShuttle() {
+      @Override public RexNode visitCall(RexCall call) {
+        return builder.cast(call, SqlTypeName.VARCHAR);
+      }
+
+      @Override public RexNode visitLocalRef(RexLocalRef localRef) {
+        if (calc.getProgram().getExprList().get(localRef.getIndex())
+            instanceof RexCall) {
+          return new RexLocalRef(localRef.getIndex(),
+              builder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR));
+        } else {
+          return localRef;
+        }
+      }
+    });
+
+    // Equivalent SQL: SELECT deptno, sal, CAST(sal + 20 AS VARCHAR) FROM emp
+    final RelNode rootWithCast =
+        builder
+            .scan("EMP")
+            .project(
+                builder.field("DEPTNO"),
+                builder.field("SAL"),
+                builder.cast(
+                    builder.call(SqlStdOperatorTable.PLUS,
+                        builder.field("SAL"), builder.literal(20)), 
SqlTypeName.VARCHAR))
+            .build();
+    assertThat(calcWithCastViaRexShuttle.getRowType(), 
is(rootWithCast.getRowType()));
+  }
 }
 
 // End RexShuttleTest.java