[CALCITE-6647] SortUnionTransposeRule should not push SORT past a UNI…

…ON when SORT's fetch is DynamicParam

core/src/main/java/org/apache/calcite/rel/rules/SortUnionTransposeRule.java

@@ -23,6 +23,7 @@
 import org.apache.calcite.rel.core.Union;
 import org.apache.calcite.rel.metadata.RelMdUtil;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexDynamicParam;
 import org.apache.calcite.tools.RelBuilderFactory;
 
 import org.immutables.value.Value;
@@ -66,11 +67,13 @@ public SortUnionTransposeRule(
   @Override public boolean matches(RelOptRuleCall call) {
     final Sort sort = call.rel(0);
     final Union union = call.rel(1);
-    // We only apply this rule if Union.all is true and Sort.offset is null.
+    // We only apply this rule if Union.all is true, Sort.offset is null and Sort.fetch is not
+    // a dynamic param.
     // There is a flag indicating if this rule should be applied when
     // Sort.fetch is null.
     return union.all
         && sort.offset == null
+        && !(sort.fetch instanceof RexDynamicParam)
         && (config.matchNullFetch() || sort.fetch != null);
   }
 

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

@@ -89,6 +89,8 @@
 import org.apache.calcite.rel.rules.ReduceExpressionsRule;
 import org.apache.calcite.rel.rules.ReduceExpressionsRule.ProjectReduceExpressionsRule;
 import org.apache.calcite.rel.rules.SingleValuesOptimizationRules;
+import org.apache.calcite.rel.rules.SortProjectTransposeRule;
+import org.apache.calcite.rel.rules.SortUnionTransposeRule;
 import org.apache.calcite.rel.rules.SpatialRules;
 import org.apache.calcite.rel.rules.UnionMergeRule;
 import org.apache.calcite.rel.rules.ValuesReduceRule;
@@ -5111,6 +5113,31 @@ private void checkEmptyJoin(RelOptFixture f) {
         .checkUnchanged();
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6647">[CALCITE-6647]
+   * SortUnionTransposeRule should not push SORT past a UNION when SORT's fetch is DynamicParam
+   </a>. */
+  @Test void testSortWithDynamicParam() {
+    HepProgramBuilder builder = new HepProgramBuilder();
+    builder.addRuleClass(SortProjectTransposeRule.class);
+    builder.addRuleClass(SortUnionTransposeRule.class);
+    HepPlanner hepPlanner = new HepPlanner(builder.build());
+    hepPlanner.addRule(CoreRules.SORT_PROJECT_TRANSPOSE);
+    hepPlanner.addRule(CoreRules.SORT_UNION_TRANSPOSE);
+    final String sql = "SELECT x.sal\n"
+        + "FROM (SELECT emp1.sal\n"
+        + "      FROM (SELECT sal\n"
+        + "            from emp\n"
+        + "            LIMIT ?) AS emp1\n"
+        + "      UNION ALL\n"
+        + "      SELECT emp2.sal\n"
+        + "      FROM (SELECT sal\n"
+        + "            from emp\n"
+        + "            LIMIT ?) AS emp2) AS x\n"
+        + "LIMIT ?";
+    sql(sql).withPlanner(hepPlanner).check();
+  }
+
   @Test void testReduceCasts() {
     // Disable simplify in RelBuilder so that there are casts in 'before';
     // The resulting plan should have no cast expressions

core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml

@@ -15623,6 +15623,51 @@ LogicalSort(sort0=[$0], dir0=[ASC], fetch=[0])
     LogicalSort(sort0=[$0], dir0=[ASC], fetch=[0])
       LogicalProject(NAME=[$1])
         LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testSortWithDynamicParam">
+    <Resource name="sql">
+      <![CDATA[SELECT x.sal
+FROM (SELECT emp1.sal
+      FROM (SELECT sal
+            from emp
+            LIMIT ?) AS emp1
+      UNION ALL
+      SELECT emp2.sal
+      FROM (SELECT sal
+            from emp
+            LIMIT ?) AS emp2) AS x
+LIMIT ?]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalSort(fetch=[?2])
+  LogicalProject(SAL=[$0])
+    LogicalUnion(all=[true])
+      LogicalProject(SAL=[$0])
+        LogicalSort(fetch=[?0])
+          LogicalProject(SAL=[$5])
+            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+      LogicalProject(SAL=[$0])
+        LogicalSort(fetch=[?1])
+          LogicalProject(SAL=[$5])
+            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+    <Resource name="planAfter">
+      <![CDATA[
+LogicalProject(SAL=[$0])
+  LogicalSort(fetch=[?2])
+    LogicalUnion(all=[true])
+      LogicalProject(SAL=[$0])
+        LogicalProject(SAL=[$5])
+          LogicalSort(fetch=[?0])
+            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+      LogicalProject(SAL=[$0])
+        LogicalProject(SAL=[$5])
+          LogicalSort(fetch=[?1])
+            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>