Disable AGG(IF()) rewrite for nondeterministic functions

yuanzhanhku · highker · commit b8e105fe68c3 · 2021-08-06T10:05:25.000-07:00
Not sure if non-deterministic functions could cause issues.
But to be safe, disable the rewrite when the IF contains any
non-deterministic functions to avoid some corner cases.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteAggregationIfToFilter.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteAggregationIfToFilter.java
@@ -29,6 +29,7 @@
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.sql.planner.iterative.Rule;
 import com.facebook.presto.sql.relational.Expressions;
+import com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedSet;
@@ -71,10 +72,12 @@ public class RewriteAggregationIfToFilter
             .with(source().matching(project().capturedAs(CHILD)));
 
     private final FunctionAndTypeManager functionAndTypeManager;
+    private final RowExpressionDeterminismEvaluator rowExpressionDeterminismEvaluator;
 
     public RewriteAggregationIfToFilter(FunctionAndTypeManager functionAndTypeManager)
     {
         this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionManager is null");
+        rowExpressionDeterminismEvaluator = new RowExpressionDeterminismEvaluator(functionAndTypeManager);
     }
 
     @Override
@@ -200,7 +203,7 @@ private boolean shouldRewriteAggregation(Aggregation aggregation, ProjectNode so
             return false;
         }
         RowExpression sourceExpression = sourceProject.getAssignments().get((VariableReferenceExpression) aggregation.getArguments().get(0));
-        if (!(sourceExpression instanceof SpecialFormExpression)) {
+        if (!(sourceExpression instanceof SpecialFormExpression) || !rowExpressionDeterminismEvaluator.isDeterministic(sourceExpression)) {
             return false;
         }
         SpecialFormExpression expression = (SpecialFormExpression) sourceExpression;
diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestRewriteAggregationIfToFilter.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestRewriteAggregationIfToFilter.java
@@ -27,6 +27,7 @@
 import java.util.Optional;
 
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.aggregation;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.expression;
@@ -72,6 +73,31 @@ public void testDoesNotFireForIfWithElse()
                 }).doesNotFire();
     }
 
+    @Test
+    public void testDoesNotFireForNonDeterministicFunction()
+    {
+        tester().assertThat(new RewriteAggregationIfToFilter(getFunctionManager()))
+                .on(p -> {
+                    VariableReferenceExpression a = p.variable("a", DOUBLE);
+                    VariableReferenceExpression ds = p.variable("ds", VARCHAR);
+                    return p.aggregation(ap -> ap.globalGrouping().step(AggregationNode.Step.FINAL)
+                            .addAggregation(p.variable("expr"), p.rowExpression("sum(a)"))
+                            .source(p.project(
+                                    assignment(a, p.rowExpression("IF(ds > '2021-07-01', random())")),
+                                    p.values(ds))));
+                }).doesNotFire();
+        tester().assertThat(new RewriteAggregationIfToFilter(getFunctionManager()))
+                .on(p -> {
+                    VariableReferenceExpression a = p.variable("a", BIGINT);
+                    VariableReferenceExpression ds = p.variable("ds", VARCHAR);
+                    return p.aggregation(ap -> ap.globalGrouping().step(AggregationNode.Step.FINAL)
+                            .addAggregation(p.variable("expr"), p.rowExpression("sum(a)"))
+                            .source(p.project(
+                                    assignment(a, p.rowExpression("IF(random() > DOUBLE '0.1', 1)")),
+                                    p.values(ds))));
+                }).doesNotFire();
+    }
+
     @Test
     public void testFireOneAggregation()
     {
