[CALCITE-4395] Add an interface in RelOptMaterializations to allow re…

core/src/main/java/org/apache/calcite/plan/RelOptMaterializations.java

@@ -22,7 +22,6 @@
 import org.apache.calcite.plan.hep.HepProgramBuilder;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.RelFactories;
-import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.sql2rel.RelFieldTrimmer;
 import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.util.Pair;
@@ -38,6 +37,7 @@
 import com.google.common.collect.Sets;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -50,6 +50,7 @@
  */
 public abstract class RelOptMaterializations {
 
+
   /**
    * Returns a list of RelNode transformed from all possible combination of
    * materialized view uses. Big queries will likely have more than one
@@ -61,23 +62,26 @@ public abstract class RelOptMaterializations {
    */
   public static List<Pair<RelNode, List<RelOptMaterialization>>> useMaterializedViews(
       final RelNode rel, List<RelOptMaterialization> materializations) {
-    return useMaterializedViews(rel, materializations, SubstitutionVisitor.DEFAULT_RULES);
+    return useMaterializedViews(rel, materializations, SubstitutionVisitor.DEFAULT_RULES,
+        SubstitutionVisitor.NORMALIZATION_RULES);
   }
 
   /**
    * Returns a list of RelNode transformed from all possible combination of
    * materialized view uses. Big queries will likely have more than one
    * transformed RelNode, e.g., (t1 group by c1) join (t2 group by c2).
-   * In addition, you can add custom materialized view recognition rules.
+   * In addition, you can add custom materialized view recognition rules and normalization rules.
    * @param rel               the original RelNode
    * @param materializations  the materialized view list
    * @param materializationRules the materialized view recognition rules
+   * @param normalizationRules  the normalization rules
    * @return the list of transformed RelNode together with their corresponding
    *         materialized views used in the transformation.
    */
   public static List<Pair<RelNode, List<RelOptMaterialization>>> useMaterializedViews(
       final RelNode rel, List<RelOptMaterialization> materializations,
-      List<SubstitutionVisitor.UnifyRule> materializationRules) {
+      List<SubstitutionVisitor.UnifyRule> materializationRules,
+      List<RelOptRule> normalizationRules) {
     final List<RelOptMaterialization> applicableMaterializations =
         getApplicableMaterializations(rel, materializations);
     final List<Pair<RelNode, List<RelOptMaterialization>>> applied =
@@ -87,7 +91,7 @@ public static List<Pair<RelNode, List<RelOptMaterialization>>> useMaterializedVi
       int count = applied.size();
       for (int i = 0; i < count; i++) {
         Pair<RelNode, List<RelOptMaterialization>> current = applied.get(i);
-        List<RelNode> sub = substitute(current.left, m, materializationRules);
+        List<RelNode> sub = substitute(current.left, m, materializationRules, normalizationRules);
         if (!sub.isEmpty()) {
           ImmutableList.Builder<RelOptMaterialization> builder =
               ImmutableList.builder();
@@ -188,7 +192,8 @@ && usesTable(materialization.qualifiedTableName, queryTablesUsed, frozenGraph))
 
   private static List<RelNode> substitute(
       RelNode root, RelOptMaterialization materialization,
-      List<SubstitutionVisitor.UnifyRule> materializationRules) {
+      List<SubstitutionVisitor.UnifyRule> materializationRules,
+      List<RelOptRule> normalizationRules) {
     // First, if the materialization is in terms of a star table, rewrite
     // the query in terms of the star table.
     if (materialization.starRelOptTable != null) {
@@ -206,20 +211,7 @@ private static List<RelNode> substitute(
     target = trimUnusedfields(target);
     HepProgram program =
         new HepProgramBuilder()
-            .addRuleInstance(CoreRules.FILTER_PROJECT_TRANSPOSE)
-            .addRuleInstance(CoreRules.FILTER_MERGE)
-            .addRuleInstance(CoreRules.FILTER_INTO_JOIN)
-            .addRuleInstance(CoreRules.JOIN_CONDITION_PUSH)
-            .addRuleInstance(CoreRules.FILTER_AGGREGATE_TRANSPOSE)
-            .addRuleInstance(CoreRules.PROJECT_MERGE)
-            .addRuleInstance(CoreRules.PROJECT_REMOVE)
-            .addRuleInstance(CoreRules.PROJECT_JOIN_TRANSPOSE)
-            .addRuleInstance(CoreRules.PROJECT_SET_OP_TRANSPOSE)
-            .addRuleInstance(CoreRules.FILTER_TO_CALC)
-            .addRuleInstance(CoreRules.PROJECT_TO_CALC)
-            .addRuleInstance(CoreRules.FILTER_CALC_MERGE)
-            .addRuleInstance(CoreRules.PROJECT_CALC_MERGE)
-            .addRuleInstance(CoreRules.CALC_MERGE)
+            .addRuleCollection(normalizationRules)
             .build();
 
     // We must use the same HEP planner for the two optimizations below.

core/src/main/java/org/apache/calcite/plan/SubstitutionVisitor.java

@@ -36,6 +36,8 @@
 import org.apache.calcite.rel.mutable.MutableScan;
 import org.apache.calcite.rel.mutable.MutableSetOp;
 import org.apache.calcite.rel.mutable.MutableUnion;
+import org.apache.calcite.rel.rules.CoreRules;
+import org.apache.calcite.rel.rules.TransformationRule;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
@@ -141,6 +143,23 @@ public class SubstitutionVisitor {
           IntersectToIntersectUnifyRule.INSTANCE,
           IntersectOnCalcsToIntersectUnifyRule.INSTANCE);
 
+  public static final ImmutableList<RelOptRule> NORMALIZATION_RULES =
+      ImmutableList.of(
+          CoreRules.FILTER_PROJECT_TRANSPOSE,
+          CoreRules.FILTER_MERGE,
+          CoreRules.FILTER_INTO_JOIN,
+          CoreRules.JOIN_CONDITION_PUSH,
+          CoreRules.FILTER_AGGREGATE_TRANSPOSE,
+          CoreRules.PROJECT_MERGE,
+          CoreRules.PROJECT_REMOVE,
+          CoreRules.PROJECT_JOIN_TRANSPOSE,
+          CoreRules.PROJECT_SET_OP_TRANSPOSE,
+          CoreRules.FILTER_TO_CALC,
+          CoreRules.PROJECT_TO_CALC,
+          CoreRules.FILTER_CALC_MERGE,
+          CoreRules.PROJECT_CALC_MERGE,
+          CoreRules.CALC_MERGE);
+
   /**
    * Factory for a builder for relational expressions.
    */

core/src/test/java/org/apache/calcite/materialize/CustomNormalizationRules.java

@@ -0,0 +1,167 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.calcite.materialize;
+
+import static org.apache.calcite.test.Matchers.isLinux;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.calcite.materialize.CustomNormalizationRules.CustomizedNormalizationRule.Config;
+import org.apache.calcite.plan.RelOptMaterialization;
+import org.apache.calcite.plan.RelOptMaterializations;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.plan.RelTraitDef;
+import org.apache.calcite.plan.SubstitutionVisitor;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalCalc;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.rules.AggregateMergeRule;
+import org.apache.calcite.rel.rules.TransformationRule;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgramBuilder;
+import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.schema.impl.AbstractTable;
+import org.apache.calcite.sql.parser.SqlParser;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.test.CalciteAssert;
+import org.apache.calcite.test.SqlToRelTestBase;
+import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.tools.RelBuilder;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.Pair;
+
+import org.junit.jupiter.api.Test;
+
+public class CustomNormalizationRules extends SqlToRelTestBase {
+
+  public static Frameworks.ConfigBuilder config() {
+    final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
+    rootSchema.add("mv0", new AbstractTable() {
+      @Override public RelDataType getRowType(RelDataTypeFactory typeFactory) {
+        return typeFactory.builder()
+            .add("deptno", SqlTypeName.INTEGER)
+            .add("count_sal", SqlTypeName.BIGINT)
+            .build();
+      }
+    });
+    return Frameworks.newConfigBuilder()
+        .parserConfig(SqlParser.Config.DEFAULT)
+        .defaultSchema(
+            CalciteAssert.addSchema(rootSchema, CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL))
+        .traitDefs((List<RelTraitDef>) null);
+  }
+
+  @Test void testCustomNormalizationRule() {
+    final RelBuilder relBuilder = RelBuilder.create(config().build());
+    final LogicalProject project = (LogicalProject) relBuilder.scan("EMP")
+        .project(relBuilder.field("EMPNO"),
+            relBuilder.field("ENAME"),
+            relBuilder.field("JOB"),
+            relBuilder.field("SAL"),
+            relBuilder.field("DEPTNO")).build();
+    final LogicalAggregate aggregate = (LogicalAggregate) relBuilder.push(project)
+        .aggregate(
+            relBuilder.groupKey(relBuilder.field(1, 0, "DEPTNO")),
+            relBuilder.count(relBuilder.field(1, 0, "SAL")))
+        .build();
+    final ImmutableBitSet groupSet = ImmutableBitSet.of(4);
+    final AggregateCall count = aggregate.getAggCallList().get(0);
+    final AggregateCall call = AggregateCall.create(count.getAggregation(),
+        count.isDistinct(), count.isApproximate(),
+        count.ignoreNulls(), ImmutableList.of(3),
+        count.filterArg, null, count.collation,
+        count.getType(), count.getName());
+    final RelNode query = LogicalAggregate.create(project, aggregate.getHints(),
+        groupSet, ImmutableList.of(groupSet), ImmutableList.of(call));
+    final RelNode target = aggregate;
+    final RelNode replacement = relBuilder.scan("mv0").build();
+    final RelOptMaterialization relOptMaterialization =
+        new RelOptMaterialization(replacement,
+            target, null, Lists.newArrayList("mv0"));
+
+    final List<RelOptRule> optRules = new ArrayList<>();
+    optRules.addAll(SubstitutionVisitor.NORMALIZATION_RULES);
+    optRules.add(CustomizedNormalizationRule.Config.DEFAULT.toRule());
+    final List<Pair<RelNode, List<RelOptMaterialization>>> relOptimized =
+        RelOptMaterializations.useMaterializedViews(query,
+            ImmutableList.of(relOptMaterialization), SubstitutionVisitor.DEFAULT_RULES, optRules);
+
+    final String optimized = ""
+        + "LogicalProject(deptno=[CAST($0):TINYINT], count_sal=[$1])\n"
+        + "  LogicalTableScan(table=[[mv0]])\n";
+    final String relOptimizedStr = RelOptUtil.toString(relOptimized.get(0).getKey());
+    assertThat(relOptimizedStr, isLinux(optimized));
+  }
+
+  /**
+   * A customized normalization rule, which compensate Calc on Agggreate.
+   */
+  public static class CustomizedNormalizationRule extends
+      RelRule<Config> implements TransformationRule {
+
+    public boolean visited = false;
+
+    public CustomizedNormalizationRule(Config config) {
+      super(config);
+    }
+
+    @Override public void onMatch(RelOptRuleCall call) {
+      Aggregate aggregate = call.rel(0);
+      RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
+      RelDataType rowType = aggregate.getRowType();
+      List<String> fieldNames = aggregate.getRowType().getFieldNames();
+      List<RexNode> projs = aggregate.getCluster().getRexBuilder()
+          .identityProjects(rowType);
+      RexProgramBuilder rexProgram = new RexProgramBuilder(rowType, rexBuilder);
+      for (int i = 0; i < projs.size(); i++) {
+        rexProgram.addProject(projs.get(i), fieldNames.get(i));
+      }
+      LogicalCalc calc = LogicalCalc
+          .create(aggregate, rexProgram.getProgram());
+      visited = true;
+      call.transformTo(calc);
+    }
+
+    /** Rule configuration. */
+    public interface Config extends RelRule.Config {
+      CustomizedNormalizationRule.Config DEFAULT = EMPTY
+          .withOperandSupplier(b ->
+              b.operand(LogicalAggregate.class).anyInputs())
+          .as(CustomizedNormalizationRule.Config.class);
+      @Override default CustomizedNormalizationRule toRule() {
+        return new CustomizedNormalizationRule(this);
+      }
+    }
+  }
+
+}

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

@@ -1691,6 +1691,12 @@ private void checkSatisfiable(RexNode e, String s) {
             + "EnumerableTableScan(table=[[hr, MV0]])")).ok();
   }
 
+  @Test public void testQueryWithoutTopCalc() {
+    final String mv = "select sum(\"empid\"), \"deptno\" from \"emps\" group by \"deptno\"";
+    final String query = "select \"deptno\" from \"emps\" group by \"deptno\"";
+    sql(mv, query).noMat();
+  }
+
   final JavaTypeFactoryImpl typeFactory =
       new JavaTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
   private final RexBuilder rexBuilder = new RexBuilder(typeFactory);