Calcite 1688

core/src/main/java/org/apache/calcite/rel/mutable/MutableMultiRel.java

@@ -16,10 +16,12 @@
  */
 package org.apache.calcite.rel.mutable;
 
+import org.apache.calcite.linq4j.Ord;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.rel.type.RelDataType;
 
 import com.google.common.base.Function;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 
 import java.util.List;
@@ -31,7 +33,11 @@ abstract class MutableMultiRel extends MutableRel {
   protected MutableMultiRel(RelOptCluster cluster,
       RelDataType rowType, MutableRelType type, List<MutableRel> inputs) {
     super(cluster, rowType, type);
-    this.inputs = inputs;
+    this.inputs = ImmutableList.copyOf(inputs);
+    for (Ord<MutableRel> input : Ord.zip(inputs)) {
+      input.e.parent = this;
+      input.e.ordinalInParent = input.i;
+    }
   }
 
   @Override public void setInput(int ordinalInParent, MutableRel input) {

core/src/test/java/org/apache/calcite/test/CalciteAssert.java

@@ -40,6 +40,8 @@
 import org.apache.calcite.util.Pair;
 import org.apache.calcite.util.Util;
 
+import org.apache.commons.lang3.StringUtils;
+
 import com.google.common.base.Function;
 import com.google.common.base.Functions;
 import com.google.common.base.Joiner;
@@ -411,6 +413,23 @@ public Void apply(ResultSet s) {
     };
   }
 
+  public static Function<ResultSet, Void> checkResultContains(
+      final String expected, final int count) {
+    return new Function<ResultSet, Void>() {
+      public Void apply(ResultSet s) {
+        try {
+          final String actual = Util.toLinux(CalciteAssert.toString(s));
+          assertTrue(
+              actual + " should have " + count + " occurrence of " + expected,
+              StringUtils.countMatches(actual, expected) == count);
+          return null;
+        } catch (SQLException e) {
+          throw new RuntimeException(e);
+        }
+      }
+    };
+  }
+
   public static Function<ResultSet, Void> checkMaskedResultContains(
       final String expected) {
     return new Function<ResultSet, Void>() {

core/src/test/java/org/apache/calcite/test/MaterializationTest.java

@@ -40,9 +40,6 @@
 import org.apache.calcite.tools.RuleSets;
 import org.apache.calcite.util.JsonBuilder;
 import org.apache.calcite.util.TryThreadLocal;
-import org.apache.calcite.util.Util;
-
-import org.apache.commons.lang3.StringUtils;
 
 import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
@@ -53,7 +50,6 @@
 
 import java.math.BigDecimal;
 import java.sql.ResultSet;
-import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -916,33 +912,24 @@ private void checkSatisfiable(RexNode e, String s) {
     checkMaterializeWithRules(m, q, rules);
   }
 
+  @Test public void testUnionAll() {
+    String q = "select * from \"emps\" where \"empid\" > 300\n"
+        + "union all select * from \"emps\" where \"empid\" < 200";
+    String m = "select * from \"emps\" where \"empid\" < 500";
+    checkMaterialize(m, q, JdbcTest.HR_MODEL,
+        CalciteAssert.checkResultContains(
+            "EnumerableTableScan(table=[[hr, m0]])", 1));
+  }
+
   @Test public void testSubQuery() {
     String q = "select \"empid\", \"deptno\", \"salary\" from \"emps\" e1\n"
         + "where \"empid\" = (\n"
         + "  select max(\"empid\") from \"emps\"\n"
         + "  where \"deptno\" = e1.\"deptno\")";
     final String m = "select \"empid\", \"deptno\" from \"emps\"\n";
-    try (final TryThreadLocal.Memo ignored = Prepare.THREAD_TRIM.push(true)) {
-      MaterializationService.setThreadLocal();
-      CalciteAssert.that()
-          .withMaterializations(JdbcTest.HR_MODEL, "m0", m)
-          .query(q)
-          .enableMaterializations(true)
-          .explainMatches("", new Function<ResultSet, Void>() {
-            public Void apply(ResultSet s) {
-              try {
-                final String actual = Util.toLinux(CalciteAssert.toString(s));
-                final String scan = "EnumerableTableScan(table=[[hr, m0]])";
-                assertTrue(actual + " should have 1 occurrence of " + scan,
-                    StringUtils.countMatches(actual, scan) == 1);
-                return null;
-              } catch (SQLException e) {
-                throw new RuntimeException(e);
-              }
-            }
-          })
-          .sameResultWithMaterializationsDisabled();
-    }
+    checkMaterialize(m, q, JdbcTest.HR_MODEL,
+        CalciteAssert.checkResultContains(
+            "EnumerableTableScan(table=[[hr, m0]])", 1));
   }
 
   @Test public void testTableModify() {
@@ -1046,28 +1033,10 @@ public List<Object> apply(JsonBuilder builder) {
     String q = "select *\n"
         + "from (select * from \"emps\" where \"empid\" < 300)\n"
         + "join (select * from \"emps\" where \"empid\" < 200) using (\"empid\")";
-    try (final TryThreadLocal.Memo ignored = Prepare.THREAD_TRIM.push(true)) {
-      MaterializationService.setThreadLocal();
-      CalciteAssert.that()
-          .withMaterializations(JdbcTest.HR_MODEL,
-              "m0", "select * from \"emps\" where \"empid\" < 500")
-          .query(q)
-          .enableMaterializations(true)
-          .explainMatches("", new Function<ResultSet, Void>() {
-            public Void apply(ResultSet s) {
-              try {
-                final String actual = Util.toLinux(CalciteAssert.toString(s));
-                final String scan = "EnumerableTableScan(table=[[hr, m0]])";
-                assertTrue(actual + " should have had two occurrences of " + scan,
-                    StringUtils.countMatches(actual, scan) == 2);
-                return null;
-              } catch (SQLException e) {
-                throw new RuntimeException(e);
-              }
-            }
-          })
-          .sameResultWithMaterializationsDisabled();
-    }
+    String m = "select * from \"emps\" where \"empid\" < 500";
+    checkMaterialize(m, q, JdbcTest.HR_MODEL,
+        CalciteAssert.checkResultContains(
+            "EnumerableTableScan(table=[[hr, m0]])", 2));
   }
 
   @Test public void testMultiMaterializationMultiUsage() {