[CALCITE-7192] AggregateReduceFunctionsRule lost FILTER condition in STDDEV/VAR function decomposition

Author: lincoln-lil

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

@@ -552,7 +552,7 @@ private static RexNode reduceStddev(
 
     final AggregateCall sumArgSquaredAggCall =
         createAggregateCallWithBinding(typeFactory, SqlStdOperatorTable.SUM,
-            argSquared.getType(), oldAggRel, oldCall, argSquaredOrdinal, -1);
+            argSquared.getType(), oldAggRel, oldCall, argSquaredOrdinal, oldCall.filterArg);
 
     final RexNode sumArgSquared =
         rexBuilder.addAggCall(sumArgSquaredAggCall,

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

@@ -2147,6 +2147,30 @@ private void checkSemiOrAntiJoinProjectTranspose(JoinRelType type) {
         .check();
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7192">[CALCITE-7192]
+   * AggregateReduceFunctionsRule lost FILTER condition in STDDEV/VAR function decomposition</a>. */
+  @Test void testVarianceStddevWithFilter() {
+    // Test to ensure FILTER conditions are correctly propagated to decomposed aggregates
+    // for all functions that use the `reduceStddev` method
+    final RelOptRule rule = AggregateReduceFunctionsRule.Config.DEFAULT
+        .withOperandFor(LogicalAggregate.class)
+        .withFunctionsToReduce(
+            EnumSet.of(SqlKind.STDDEV_POP, SqlKind.STDDEV_SAMP,
+                      SqlKind.VAR_POP, SqlKind.VAR_SAMP, SqlKind.AVG))
+        .toRule();
+    final String sql = "select name, "
+        + "stddev_pop(deptno) filter (where deptno > 10), "
+        + "stddev_samp(deptno) filter (where deptno > 20), "
+        + "var_pop(deptno) filter (where deptno > 30), "
+        + "var_samp(deptno) filter (where deptno > 40), "
+        + "avg(deptno) filter (where deptno > 50)\n"
+        + "from sales.dept group by name";
+    sql(sql)
+        .withRule(rule)
+        .check();
+  }
+
   @Test void testDistinctCountWithoutGroupBy() {
     final String sql = "select max(deptno), count(distinct ename)\n"
         + "from sales.emp";

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

@@ -20068,6 +20068,28 @@ LogicalProject(DEPTNO=[$7])
       <![CDATA[
 LogicalProject(DEPTNO=[$0], NAME=[$1])
   LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testVarianceStddevWithFilter">
+    <Resource name="sql">
+      <![CDATA[select name, stddev_pop(deptno) filter (where deptno > 10), stddev_samp(deptno) filter (where deptno > 20), var_pop(deptno) filter (where deptno > 30), var_samp(deptno) filter (where deptno > 40), avg(deptno) filter (where deptno > 50)
+from sales.dept group by name]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalAggregate(group=[{0}], EXPR$1=[STDDEV_POP($1) FILTER $2], EXPR$2=[STDDEV_SAMP($1) FILTER $3], EXPR$3=[VAR_POP($1) FILTER $4], EXPR$4=[VAR_SAMP($1) FILTER $5], EXPR$5=[AVG($1) FILTER $6])
+  LogicalProject(NAME=[$1], DEPTNO=[$0], $f2=[>($0, 10)], $f3=[>($0, 20)], $f4=[>($0, 30)], $f5=[>($0, 40)], $f6=[>($0, 50)])
+    LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+    </Resource>
+    <Resource name="planAfter">
+      <![CDATA[
+LogicalProject(NAME=[$0], EXPR$1=[CAST(POWER(/(-($1, /(*($2, $2), $3)), $3), 0.5:DECIMAL(2, 1))):INTEGER], EXPR$2=[CAST(POWER(/(-($4, /(*($5, $5), $6)), CASE(=($6, 1), null:BIGINT, -($6, 1))), 0.5:DECIMAL(2, 1))):INTEGER], EXPR$3=[CAST(/(-($7, /(*($8, $8), $9)), $9)):INTEGER], EXPR$4=[CAST(/(-($10, /(*($11, $11), $12)), CASE(=($12, 1), null:BIGINT, -($12, 1)))):INTEGER], EXPR$5=[CAST(/($13, $14)):INTEGER])
+  LogicalAggregate(group=[{0}], agg#0=[SUM($7) FILTER $2], agg#1=[SUM($1) FILTER $2], agg#2=[COUNT() FILTER $2], agg#3=[SUM($7) FILTER $3], agg#4=[SUM($1) FILTER $3], agg#5=[COUNT() FILTER $3], agg#6=[SUM($7) FILTER $4], agg#7=[SUM($1) FILTER $4], agg#8=[COUNT() FILTER $4], agg#9=[SUM($7) FILTER $5], agg#10=[SUM($1) FILTER $5], agg#11=[COUNT() FILTER $5], agg#12=[SUM($1) FILTER $6], agg#13=[COUNT() FILTER $6])
+    LogicalProject(NAME=[$0], DEPTNO=[$1], $f2=[$2], $f3=[$3], $f4=[$4], $f5=[$5], $f6=[$6], $f7=[*($1, $1)])
+      LogicalProject(NAME=[$1], DEPTNO=[$0], $f2=[>($0, 10)], $f3=[>($0, 20)], $f4=[>($0, 30)], $f5=[>($0, 40)], $f6=[>($0, 50)])
+        LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
 ]]>
     </Resource>
   </TestCase>

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

@@ -239,6 +239,30 @@ group by gender;
 
 !ok
 
+# [CALCITE-7192] lost FILTER condition when decomposition
+select gender,
+  stddev_pop(deptno) as p,
+  stddev_pop(deptno) filter (where deptno > 50) as pf,
+  stddev_samp(deptno) as s,
+  stddev_samp(deptno) filter (where deptno > 40) as sf,
+  var_pop(deptno) as vp,
+  var_pop(deptno) filter (where deptno > 30) as vpf,
+  var_samp(deptno) as vs,
+  var_samp(deptno) filter (where deptno > 20) as vsf,
+  avg(deptno) as a,
+  avg(deptno) filter (where deptno > 10) as af
+from emp
+group by gender;
++--------+----+----+----+----+-----+-----+-----+-----+----+----+
+| GENDER | P  | PF | S  | SF | VP  | VPF | VS  | VSF | A  | AF |
++--------+----+----+----+----+-----+-----+-----+-----+----+----+
+| F      | 17 |  0 | 19 |  7 | 304 |  25 | 380 | 225 | 36 | 42 |
+| M      | 17 |    | 20 |    | 289 |   0 | 433 |     | 26 | 35 |
++--------+----+----+----+----+-----+-----+-----+-----+----+----+
+(2 rows)
+
+!ok
+
 select city, gender as c from emps;
 +---------------+---+
 | CITY          | C |