This is an automated email from the ASF dual-hosted git repository.
jhyde pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git
The following commit(s) were added to refs/heads/master by this push:
new ecbafbf [CALCITE-4925] AggregateReduceFunctionsRule should accept
arbitrary predicates
ecbafbf is described below
commit ecbafbfbd356955db08e363b71c03d8923e6fe98
Author: Will Noble <[email protected]>
AuthorDate: Wed Dec 1 02:38:25 2021 -0800
[CALCITE-4925] AggregateReduceFunctionsRule should accept arbitrary
predicates
Allow selective application of the reduce-functions rule with arbitrary
specificity.
Close apache/calcite#2634
---
.../rel/rules/AggregateReduceFunctionsRule.java | 33 ++++++++++++++++++----
.../org/apache/calcite/test/RelOptRulesTest.java | 10 +++++++
.../org/apache/calcite/test/RelOptRulesTest.xml | 22 +++++++++++++++
3 files changed, 59 insertions(+), 6 deletions(-)
diff --git
a/core/src/main/java/org/apache/calcite/rel/rules/AggregateReduceFunctionsRule.java
b/core/src/main/java/org/apache/calcite/rel/rules/AggregateReduceFunctionsRule.java
index 47dd5d1..522835f 100644
---
a/core/src/main/java/org/apache/calcite/rel/rules/AggregateReduceFunctionsRule.java
+++
b/core/src/main/java/org/apache/calcite/rel/rules/AggregateReduceFunctionsRule.java
@@ -53,6 +53,7 @@ import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.IntPredicate;
+import java.util.function.Predicate;
/**
* Planner rule that reduces aggregate functions in
@@ -173,7 +174,8 @@ public class AggregateReduceFunctionsRule
/** Returns whether this rule can reduce a given aggregate function call. */
public boolean canReduce(AggregateCall call) {
- return functionsToReduce.contains(call.getAggregation().getKind());
+ return functionsToReduce.contains(call.getAggregation().getKind())
+ && config.extraCondition().test(call);
}
/**
@@ -209,7 +211,7 @@ public class AggregateReduceFunctionsRule
relBuilder.push(oldAggRel.getInput());
final List<RexNode> inputExprs = new ArrayList<>(relBuilder.fields());
- // create new agg function calls and rest of project list together
+ // create new aggregate function calls and rest of project list together
for (AggregateCall oldCall : oldCalls) {
projList.add(
reduceAgg(
@@ -798,7 +800,7 @@ public class AggregateReduceFunctionsRule
}
/**
- * Do a shallow clone of oldAggRel and update aggCalls. Could be refactored
+ * Does a shallow clone of oldAggRel and updates aggCalls. Could be
refactored
* into Aggregate and subclasses - but it's only needed for some
* subclasses.
*
@@ -816,13 +818,13 @@ public class AggregateReduceFunctionsRule
}
/**
- * Add a calc with the expressions to compute the original agg calls from the
- * decomposed ones.
+ * Adds a calculation with the expressions to compute the original aggregate
+ * calls from the decomposed ones.
*
* @param relBuilder Builder of relational expressions; at the top of its
* stack is its input
* @param rowType The output row type of the original aggregate.
- * @param exprs The expressions to compute the original agg calls.
+ * @param exprs The expressions to compute the original aggregate calls
*/
protected void newCalcRel(RelBuilder relBuilder,
RelDataType rowType,
@@ -847,8 +849,24 @@ public class AggregateReduceFunctionsRule
return new AggregateReduceFunctionsRule(this);
}
+ /** The set of aggregate function types to try to reduce.
+ *
+ * <p>Any aggregate function whose type is omitted from this set, OR which
+ * does not pass the {@link #extraCondition}, will be ignored.
+ */
@Nullable Set<SqlKind> functionsToReduce();
+ /** A test that must pass before attempting to reduce any aggregate
function.
+ *
+ * <p>Any aggegate function which does not pass, OR whose type is omitted
+ * from {@link #functionsToReduce}, will be ignored. The default predicate
+ * always passes.
+ */
+ @Value.Default
+ default Predicate<AggregateCall> extraCondition() {
+ return ignored -> true;
+ }
+
/** Sets {@link #functionsToReduce}. */
Config withFunctionsToReduce(@Nullable Iterable<SqlKind> functionSet);
@@ -856,6 +874,9 @@ public class AggregateReduceFunctionsRule
return withFunctionsToReduce((Iterable<SqlKind>) functionSet);
}
+ /** Sets {@link #extraCondition}. */
+ Config withExtraCondition(Predicate<AggregateCall> test);
+
/** Returns the validated set of functions to reduce, or the default set
* if not specified. */
default Set<SqlKind> actualFunctionsToReduce() {
diff --git a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
index 2cad6a9..a01731a 100644
--- a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
@@ -6321,6 +6321,16 @@ class RelOptRulesTest extends RelOptTestBase {
sql(sql).withRule(rule).check();
}
+ @Test void testReduceWithNonTypePredicate() {
+ // Make sure we can reduce with more specificity than just agg function
type.
+ final RelOptRule rule = AggregateReduceFunctionsRule.Config.DEFAULT
+ .withExtraCondition(call -> call.distinctKeys != null)
+ .toRule();
+ final String sql = "select avg(sal), avg(sal) within distinct (deptno)\n"
+ + "from emp";
+ sql(sql).withRule(rule).check();
+ }
+
/** Test case for
* <a href="https://issues.apache.org/jira/browse/CALCITE-2803";>[CALCITE-2803]
* Identify expanded IS NOT DISTINCT FROM expression when pushing project
past join</a>.
diff --git
a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
index 2b22095..e9a4152 100644
--- a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
+++ b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
@@ -10877,6 +10877,28 @@ LogicalValues(tuples=[[{ 11, 1, 10 }, { 23, 3, 20 }]])
]]>
</Resource>
</TestCase>
+ <TestCase name="testReduceWithNonTypePredicate">
+ <Resource name="sql">
+ <![CDATA[select avg(sal), avg(sal) within distinct (deptno)
+from emp]]>
+ </Resource>
+ <Resource name="planBefore">
+ <![CDATA[
+LogicalAggregate(group=[{}], EXPR$0=[AVG($0)], EXPR$1=[AVG($0) WITHIN DISTINCT
($1)])
+ LogicalProject(SAL=[$5], DEPTNO=[$7])
+ LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+ </Resource>
+ <Resource name="planAfter">
+ <![CDATA[
+LogicalProject(EXPR$0=[$0], EXPR$1=[CAST(/($1, $2)):INTEGER])
+ LogicalProject(EXPR$0=[$0], $f1=[CASE(=($2, 0), null:INTEGER, $1)], $f2=[$2])
+ LogicalAggregate(group=[{}], EXPR$0=[AVG($0)], agg#1=[$SUM0($0) WITHIN
DISTINCT ($1)], agg#2=[COUNT() WITHIN DISTINCT ($1)])
+ LogicalProject(SAL=[$5], DEPTNO=[$7])
+ LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+ </Resource>
+ </TestCase>
<TestCase name="testRemoveDistinctOnAgg">
<Resource name="sql">
<