[CALCITE-6335] Quidem tests should allow specifying optimization passes to apply to programs

Author: xiedeyantu

core/src/test/resources/sql/planner.iq

@@ -0,0 +1,89 @@
+# planner.iq - planner tests can customizable optimization rules
+#
+# 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.
+#
+!use post
+!set outputformat mysql
+
+# The following line is used to update the default planner rules.
+# "+" means add a rule, "-" means remove a rule. Only the "CoreRules." can be omitted.
+!set planner-rules "-INTERSECT_TO_DISTINCT,-EnumerableRules.ENUMERABLE_INTERSECT_RULE,+CoreRules.INTERSECT_TO_SEMI_JOIN"
+
+# Test INTERSECT_TO_SEMI_JOIN
+with t (i) as (values (0), (1))
+select * from t as t1
+intersect
+select * from t as t2 where t2.i > 0;
++---+
+| I |
++---+
+| 1 |
++---+
+(1 row)
+
+!ok
+
+EnumerableNestedLoopJoin(condition=[IS NOT DISTINCT FROM($0, $1)], joinType=[semi])
+  EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+  EnumerableCalc(expr#0=[{inputs}], expr#1=[0], expr#2=[>($t0, $t1)], EXPR$0=[$t0], $condition=[$t2])
+    EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+!plan
+
+# Test that rules set by "!set planner-rules" are sticky
+with t (i) as (values (0), (1))
+select * from t as t1
+intersect
+select * from t as t2 where t2.i > 0;
++---+
+| I |
++---+
+| 1 |
++---+
+(1 row)
+
+!ok
+
+EnumerableNestedLoopJoin(condition=[IS NOT DISTINCT FROM($0, $1)], joinType=[semi])
+  EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+  EnumerableCalc(expr#0=[{inputs}], expr#1=[0], expr#2=[>($t0, $t1)], EXPR$0=[$t0], $condition=[$t2])
+    EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+!plan
+
+# This command is used to reset the planner rules to the original state.
+# Before running this command, the planner rules by "!set planner-rules" are sticky.
+!set planner-rules original
+
+# Test that rules set by "!set planner-rules original" are reset
+with t (i) as (values (0), (1))
+select * from t as t1
+intersect
+select * from t as t2 where t2.i > 0;
++---+
+| I |
++---+
+| 1 |
++---+
+(1 row)
+
+!ok
+
+EnumerableIntersect(all=[false])
+  EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+  EnumerableCalc(expr#0=[{inputs}], expr#1=[0], expr#2=[>($t0, $t1)], EXPR$0=[$t0], $condition=[$t2])
+    EnumerableValues(tuples=[[{ 0 }, { 1 }]])
+!plan
+
+# End planner.iq

testkit/src/main/java/org/apache/calcite/test/QuidemTest.java

@@ -16,10 +16,14 @@
  */
 package org.apache.calcite.test;
 
+import org.apache.calcite.adapter.enumerable.EnumerableRules;
 import org.apache.calcite.avatica.AvaticaUtils;
 import org.apache.calcite.config.CalciteConnectionProperty;
 import org.apache.calcite.jdbc.CalciteConnection;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.prepare.Prepare;
+import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.runtime.Hook;
@@ -56,7 +60,9 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static org.junit.jupiter.api.Assertions.fail;
@@ -71,6 +77,9 @@ public abstract class QuidemTest {
 
   private static final Pattern PATTERN = Pattern.compile("\\.iq$");
 
+  // Saved original planner rules
+  private static @Nullable List<RelOptRule> originalRules;
+
   private static @Nullable Object getEnv(String varName) {
     switch (varName) {
     case "jdk18":
@@ -171,6 +180,26 @@ protected void checkRun(String path) throws Exception {
               int thresholdValue = ((BigDecimal) value).intValue();
               closer.add(Prepare.THREAD_INSUBQUERY_THRESHOLD.push(thresholdValue));
             }
+            // Configures query planner rules via "!set planner-rules" command.
+            // The value can be set as follows:
+            // - Add rule:       "+EnumerableRules.ENUMERABLE_INTERSECT_RULE"
+            // - Remove rule:    "-CoreRules.INTERSECT_TO_DISTINCT"
+            // - Short form:     "+INTERSECT_TO_DISTINCT" (CoreRules prefix may be omitted)
+            // - Reset defaults: "original"
+            if (propertyName.equals("planner-rules")) {
+              if (value.equals("original")) {
+                closer.add(Hook.PLANNER.addThread(this::resetPlanner));
+              } else {
+                closer.add(
+                    Hook.PLANNER.addThread((Consumer<RelOptPlanner>)
+                        planner -> {
+                          if (originalRules == null) {
+                            originalRules = planner.getRules();
+                          }
+                          updatePlanner(planner, (String) value);
+                        }));
+              }
+            }
           })
           .withEnv(QuidemTest::getEnv)
           .build();
@@ -187,6 +216,57 @@ protected void checkRun(String path) throws Exception {
     }
   }
 
+  private void updatePlanner(RelOptPlanner planner, String value) {
+    List<RelOptRule> rulesAdd = new ArrayList<>();
+    List<RelOptRule> rulesRemove = new ArrayList<>();
+    parseRules(value, rulesAdd, rulesRemove);
+    rulesRemove.forEach(planner::removeRule);
+    rulesAdd.forEach(planner::addRule);
+  }
+
+  private void resetPlanner(RelOptPlanner planner) {
+    if (originalRules != null) {
+      planner.getRules().forEach(planner::removeRule);
+      originalRules.forEach(planner::addRule);
+    }
+  }
+
+  private void parseRules(String value, List<RelOptRule> rulesAdd, List<RelOptRule> rulesRemove) {
+    Pattern pattern = Pattern.compile("([+-])((CoreRules|EnumerableRules)\\.)?(\\w+)");
+    Matcher matcher = pattern.matcher(value);
+
+    while (matcher.find()) {
+      char operation = matcher.group(1).charAt(0);
+      String ruleSource = matcher.group(3);
+      String ruleName = matcher.group(4);
+
+      try {
+        if (ruleSource == null || ruleSource.equals("CoreRules")) {
+          Object rule = CoreRules.class.getField(ruleName).get(null);
+          setRules(operation, (RelOptRule) rule, rulesAdd, rulesRemove);
+        } else if (ruleSource.equals("EnumerableRules")) {
+          Object rule = EnumerableRules.class.getField(ruleName).get(null);
+          setRules(operation, (RelOptRule) rule, rulesAdd, rulesRemove);
+        } else {
+          throw new RuntimeException("Unknown rule: " + ruleName);
+        }
+      } catch (NoSuchFieldException | IllegalAccessException e) {
+        throw new RuntimeException("set rules failed: " + e.getMessage(), e);
+      }
+    }
+  }
+
+  private void setRules(char operation, RelOptRule rule,
+      List<RelOptRule> rulesAdd, List<RelOptRule> rulesRemove) {
+    if (operation == '+') {
+      rulesAdd.add(rule);
+    } else if (operation == '-') {
+      rulesRemove.add(rule);
+    } else {
+      throw new RuntimeException("unknown operation '" + operation + "'");
+    }
+  }
+
   /** Returns a file, replacing one directory with another.
    *
    * <p>For example, {@code replaceDir("/abc/str/astro.txt", "str", "xyz")}