[CALCITE-7779] Optimization rule FilterAggregateTransposeRule rewrites queries to semantically non-equivalent ones

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

@@ -32,6 +32,7 @@
 import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.tools.RelBuilderFactory;
 import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.mapping.Mappings;
 
 import com.google.common.collect.ImmutableList;
 
@@ -149,8 +150,13 @@ private static boolean canPush(Aggregate aggregate, ImmutableBitSet rCols) {
       // If grouping sets are used, the filter can be pushed if
       // the columns referenced in the predicate are present in
       // all the grouping sets.
+      final Mappings.TargetMapping aggregateInputToOutput =
+          Mappings.target(aggregate.getGroupSet().asList(),
+              aggregate.getInput().getRowType().getFieldCount());
       for (ImmutableBitSet groupingSet : aggregate.getGroupSets()) {
-        if (!groupingSet.contains(rCols)) {
+        final ImmutableBitSet groupingSetOutputRefs =
+            groupingSet.permute(aggregateInputToOutput);
+        if (!groupingSetOutputRefs.contains(rCols)) {
           return false;
         }
       }

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

@@ -1364,6 +1364,32 @@ private RelOptFixture basePushFilterPastAggWithGroupingSets() {
     basePushFilterPastAggWithGroupingSets().check();
   }
 
+  private RelOptFixture pushFilterPastAggWithNonLeadingGroupingKeys(String condition) {
+    final String sql = "select a, b\n"
+        + "from (values (0, 1, 2)) as t(unused, a, b)\n"
+        + "group by grouping sets ((a), (a, b))\n"
+        + "having " + condition;
+    return sql(sql)
+        .withPreRule(CoreRules.AGGREGATE_PROJECT_MERGE)
+        .withRule(CoreRules.FILTER_AGGREGATE_TRANSPOSE);
+  }
+
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7779">[CALCITE-7779]
+   * Optimization rule FilterAggregateTransposeRule rewrites queries to
+   * semantically non-equivalent ones</a>. */
+  @Test void testPushFilterPastAggWithGroupingSetsNonLeadingKeysIsNull() {
+    pushFilterPastAggWithNonLeadingGroupingKeys("b is null").checkUnchanged();
+  }
+
+  @Test void testPushFilterPastAggWithGroupingSetsNonLeadingKeysEquals() {
+    pushFilterPastAggWithNonLeadingGroupingKeys("b = 2").checkUnchanged();
+  }
+
+  @Test void testPushFilterPastAggWithGroupingSetsNonLeadingKeys() {
+    pushFilterPastAggWithNonLeadingGroupingKeys("a = 1").check();
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-434">[CALCITE-434]
    * FilterAggregateTransposeRule loses conditions that cannot be pushed</a>. */

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

@@ -15104,6 +15104,58 @@ LogicalProject(DNAME=[$0], DDEPTNO=[$1], C=[$2])
     LogicalFilter(condition=[=($0, 'Charlie')])
       LogicalProject(DNAME=[$1], DDEPTNO=[$0])
         LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testPushFilterPastAggWithGroupingSetsNonLeadingKeys">
+    <Resource name="sql">
+      <![CDATA[select a, b
+from (values (0, 1, 2)) as t(unused, a, b)
+group by grouping sets ((a), (a, b))
+having a = 1]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalFilter(condition=[=($0, 1)])
+  LogicalAggregate(group=[{1, 2}], groups=[[{1, 2}, {1}]])
+    LogicalValues(tuples=[[{ 0, 1, 2 }]])
+]]>
+    </Resource>
+    <Resource name="planAfter">
+      <![CDATA[
+LogicalAggregate(group=[{1, 2}], groups=[[{1, 2}, {1}]])
+  LogicalFilter(condition=[=($1, 1)])
+    LogicalValues(tuples=[[{ 0, 1, 2 }]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testPushFilterPastAggWithGroupingSetsNonLeadingKeysEquals">
+    <Resource name="sql">
+      <![CDATA[select a, b
+from (values (0, 1, 2)) as t(unused, a, b)
+group by grouping sets ((a), (a, b))
+having b = 2]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalFilter(condition=[=($1, 2)])
+  LogicalAggregate(group=[{1, 2}], groups=[[{1, 2}, {1}]])
+    LogicalValues(tuples=[[{ 0, 1, 2 }]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testPushFilterPastAggWithGroupingSetsNonLeadingKeysIsNull">
+    <Resource name="sql">
+      <![CDATA[select a, b
+from (values (0, 1, 2)) as t(unused, a, b)
+group by grouping sets ((a), (a, b))
+having b is null]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalFilter(condition=[IS NULL($1)])
+  LogicalAggregate(group=[{1, 2}], groups=[[{1, 2}, {1}]])
+    LogicalValues(tuples=[[{ 0, 1, 2 }]])
 ]]>
     </Resource>
   </TestCase>