[CALCITE-3848] Rewriting for materialized view consisting of group by…

… on join keys fails with Mappings$NoElementException (Vineet Garg)

Close apache#1852

core/src/main/java/org/apache/calcite/rel/rules/materialize/MaterializedViewAggregateRule.java

@@ -63,6 +63,7 @@
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Multimap;
 
@@ -366,7 +367,6 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
           rexBuilder.makeInputRef(relBuilder.peek(),
               aggregate.getGroupCount() + i);
       aggregateCalls.add(
-          // TODO: handle aggregate ordering
           relBuilder.aggregateCall(rollupAgg, operand)
               .distinct(aggCall.isDistinct())
               .approximate(aggCall.isApproximate())
@@ -528,7 +528,10 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
 
     // Generate result rewriting
     final List<RexNode> additionalViewExprs = new ArrayList<>();
-    Mapping rewritingMapping = null;
+
+    // Multimap is required since a column in the materialized view's project
+    // could map to multiple columns in the target query
+    ImmutableMultimap<Integer, Integer> rewritingMapping = null;
     RelNode result = relBuilder.push(input).build();
     // We create view expressions that will be used in a Project on top of the
     // view in case we need to rollup the expression
@@ -542,9 +545,8 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
         return null;
       }
       // Target is coarser level of aggregation. Generate an aggregate.
-      rewritingMapping = Mappings.create(MappingType.FUNCTION,
-          topViewProject.getRowType().getFieldCount() + viewAggregateAdditionalFieldCount,
-          queryAggregate.getRowType().getFieldCount());
+      final ImmutableMultimap.Builder<Integer, Integer> rewritingMappingB =
+          ImmutableMultimap.builder();
       final ImmutableBitSet.Builder groupSetB = ImmutableBitSet.builder();
       for (int i = 0; i < queryAggregate.getGroupCount(); i++) {
         int targetIdx = aggregateMapping.getTargetOpt(i);
@@ -580,7 +582,7 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
           }
           // We create the new node pointing to the index
           groupSetB.set(inputViewExprs.size());
-          rewritingMapping.set(inputViewExprs.size(), i);
+          rewritingMappingB.put(inputViewExprs.size(), i);
           additionalViewExprs.add(
               new RexInputRef(targetIdx, targetNode.getType()));
           // We need to create the rollup expression
@@ -597,7 +599,7 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
             int ref = ((RexInputRef) n).getIndex();
             if (ref == targetIdx) {
               groupSetB.set(k);
-              rewritingMapping.set(k, i);
+              rewritingMappingB.put(k, i);
               added = true;
             }
           }
@@ -636,10 +638,9 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
               // Cannot rollup this aggregate, bail out
               return null;
             }
-            rewritingMapping.set(k, queryAggregate.getGroupCount() + aggregateCalls.size());
+            rewritingMappingB.put(k, queryAggregate.getGroupCount() + aggregateCalls.size());
             final RexInputRef operand = rexBuilder.makeInputRef(input, k);
             aggregateCalls.add(
-                // TODO: handle aggregate ordering
                 relBuilder.aggregateCall(rollupAgg, operand)
                     .approximate(queryAggCall.isApproximate())
                     .distinct(queryAggCall.isDistinct())
@@ -669,12 +670,13 @@ protected MaterializedViewAggregateRule(RelOptRuleOperand operand,
             .build();
       }
       // We introduce a project on top, as group by columns order is lost
-      List<RexNode> projects = new ArrayList<>();
-      Mapping inverseMapping = rewritingMapping.inverse();
+      rewritingMapping = rewritingMappingB.build();
+      final ImmutableMultimap<Integer, Integer> inverseMapping = rewritingMapping.inverse();
+      final List<RexNode> projects = new ArrayList<>();
       for (int i = 0; i < queryAggregate.getGroupCount(); i++) {
         projects.add(
             rexBuilder.makeInputRef(result,
-                groupSet.indexOf(inverseMapping.getTarget(i))));
+                groupSet.indexOf(inverseMapping.get(i).iterator().next())));
       }
       // We add aggregate functions that are present in result to projection list
       for (int i = queryAggregate.getGroupCount(); i < result.getRowType().getFieldCount(); i++) {

core/src/main/java/org/apache/calcite/rel/rules/materialize/MaterializedViewRule.java

@@ -1175,7 +1175,7 @@ protected RexNode shuttleReferences(final RexBuilder rexBuilder,
    */
   protected RexNode shuttleReferences(final RexBuilder rexBuilder,
       final RexNode expr, final Multimap<RexNode, Integer> exprsLineage,
-      final RelNode node, final Mapping rewritingMapping) {
+      final RelNode node, final Multimap<Integer, Integer> rewritingMapping) {
     try {
       RexShuttle visitor =
           new RexShuttle() {
@@ -1187,11 +1187,11 @@ protected RexNode shuttleReferences(final RexBuilder rexBuilder,
               }
               int pos = c.iterator().next();
               if (rewritingMapping != null) {
-                pos = rewritingMapping.getTargetOpt(pos);
-                if (pos == -1) {
+                if (!rewritingMapping.containsKey(pos)) {
                   // Cannot map expression
                   throw Util.FoundOne.NULL;
                 }
+                pos = rewritingMapping.get(pos).iterator().next();
               }
               if (node != null) {
                 return rexBuilder.makeInputRef(node, pos);
@@ -1207,11 +1207,11 @@ protected RexNode shuttleReferences(final RexBuilder rexBuilder,
               }
               int pos = c.iterator().next();
               if (rewritingMapping != null) {
-                pos = rewritingMapping.getTargetOpt(pos);
-                if (pos == -1) {
+                if (!rewritingMapping.containsKey(pos)) {
                   // Cannot map expression
                   throw Util.FoundOne.NULL;
                 }
+                pos = rewritingMapping.get(pos).iterator().next();
               }
               if (node != null) {
                 return rexBuilder.makeInputRef(node, pos);
@@ -1227,11 +1227,11 @@ protected RexNode shuttleReferences(final RexBuilder rexBuilder,
               }
               int pos = c.iterator().next();
               if (rewritingMapping != null) {
-                pos = rewritingMapping.getTargetOpt(pos);
-                if (pos == -1) {
+                if (!rewritingMapping.containsKey(pos)) {
                   // Cannot map expression
                   return super.visitCall(call);
                 }
+                pos = rewritingMapping.get(pos).iterator().next();
               }
               if (node != null) {
                 return rexBuilder.makeInputRef(node, pos);

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

@@ -2511,6 +2511,21 @@ private void checkSatisfiable(RexNode e, String s) {
         .ok();
   }
 
+  @Test public void testAggregateOnJoinKeys() {
+    sql("select \"deptno\", \"empid\", \"salary\" "
+            + "from \"emps\"\n"
+            + "group by \"deptno\", \"empid\", \"salary\"",
+        "select \"empid\", \"depts\".\"deptno\" "
+            + "from \"emps\"\n"
+            + "join \"depts\" on \"depts\".\"deptno\" = \"empid\" group by \"empid\", \"depts\".\"deptno\"")
+        .withResultContains(
+            "EnumerableCalc(expr#0=[{inputs}], empid=[$t0], empid0=[$t0])\n"
+                + "  EnumerableAggregate(group=[{1}])\n"
+                + "    EnumerableHashJoin(condition=[=($1, $3)], joinType=[inner])\n"
+                + "      EnumerableTableScan(table=[[hr, m0]])")
+        .ok();
+  }
+
   @Test public void testViewMaterialization() {
     sql("select \"depts\".\"name\"\n"
             + "from \"emps\"\n"