Optimize filter condition with CASE predicate

prestodb · May 12, 2022 · 67c8e6d · 67c8e6d
1 parent b17ec68
commit 67c8e6d
Show file tree

Hide file tree

Showing 10 changed files with 645 additions and 1 deletion.
diff --git a/presto-expressions/src/main/java/com/facebook/presto/expressions/LogicalRowExpressions.java b/presto-expressions/src/main/java/com/facebook/presto/expressions/LogicalRowExpressions.java
@@ -51,6 +51,7 @@
 import static com.facebook.presto.spi.relation.SpecialFormExpression.Form.AND;
 import static com.facebook.presto.spi.relation.SpecialFormExpression.Form.IS_NULL;
 import static com.facebook.presto.spi.relation.SpecialFormExpression.Form.OR;
+import static com.facebook.presto.spi.relation.SpecialFormExpression.Form.SWITCH;
 import static java.lang.Math.min;
 import static java.util.Arrays.asList;
 import static java.util.Arrays.stream;
@@ -631,6 +632,21 @@ private boolean isComparisonExpression(RowExpression expression)
         return expression instanceof CallExpression && functionResolution.isComparisonFunction(((CallExpression) expression).getFunctionHandle());
     }
 
+    public boolean isEqualsExpression(RowExpression expression)
+    {
+        return expression instanceof CallExpression && functionResolution.isEqualsFunction(((CallExpression) expression).getFunctionHandle());
+    }
+
+    public boolean isCastExpression(RowExpression expression)
+    {
+        return expression instanceof CallExpression && functionResolution.isCastFunction(((CallExpression) expression).getFunctionHandle());
+    }
+
+    public boolean isCaseExpression(RowExpression expression)
+    {
+        return expression instanceof SpecialFormExpression && ((SpecialFormExpression) expression).getForm().equals(SWITCH);
+    }
+
     /**
      * Extract the component predicates as a list of list in which is grouped so that the outer level has same conjunctive/disjunctive joiner as original predicate and
      * inner level has opposite joiner.
@@ -771,11 +787,29 @@ private Optional<OperatorType> getOperator(RowExpression expression)
         return Optional.empty();
     }
 
-    private RowExpression notCallExpression(RowExpression argument)
+    public RowExpression notCallExpression(RowExpression argument)
     {
         return new CallExpression(argument.getSourceLocation(), "not", functionResolution.notFunction(), BOOLEAN, singletonList(argument));
     }
 
+    public RowExpression equalsCallExpression(RowExpression left, RowExpression right)
+    {
+        return new CallExpression(
+                EQUAL.name(),
+                functionResolution.comparisonFunction(EQUAL, left.getType(), right.getType()),
+                BOOLEAN,
+                asList(left, right));
+    }
+
+    public static RowExpression replaceArguments(CallExpression expression, RowExpression... arguments)
+    {
+        return new CallExpression(
+                expression.getDisplayName(),
+                expression.getFunctionHandle(),
+                expression.getType(),
+                asList(arguments));
+    }
+
     private static OperatorType negate(OperatorType operator)
     {
         switch (operator) {

diff --git a/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java b/presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java
@@ -159,6 +159,7 @@ public final class SystemSessionProperties
     public static final String PARTIAL_AGGREGATION_STRATEGY = "partial_aggregation_strategy";
     public static final String PARTIAL_AGGREGATION_BYTE_REDUCTION_THRESHOLD = "partial_aggregation_byte_reduction_threshold";
     public static final String OPTIMIZE_TOP_N_ROW_NUMBER = "optimize_top_n_row_number";
+    public static final String OPTIMIZE_CASE_EXPRESSION_PREDICATE = "optimize_case_expression_predicate";
     public static final String MAX_GROUPING_SETS = "max_grouping_sets";
     public static final String LEGACY_UNNEST = "legacy_unnest";
     public static final String STATISTICS_CPU_TIMER_ENABLED = "statistics_cpu_timer_enabled";
@@ -851,6 +852,11 @@ public SystemSessionProperties(
                         "Use top N row number optimization",
                         featuresConfig.isOptimizeTopNRowNumber(),
                         false),
+                booleanProperty(
+                        OPTIMIZE_CASE_EXPRESSION_PREDICATE,
+                        "Optimize case expression predicates",
+                        featuresConfig.isOptimizeCaseExpressionPredicate(),
+                        false),
                 integerProperty(
                         MAX_GROUPING_SETS,
                         "Maximum number of grouping sets in a GROUP BY",
@@ -1795,6 +1801,11 @@ public static boolean isOptimizeTopNRowNumber(Session session)
         return session.getSystemProperty(OPTIMIZE_TOP_N_ROW_NUMBER, Boolean.class);
     }
 
+    public static boolean isOptimizeCaseExpressionPredicate(Session session)
+    {
+        return session.getSystemProperty(OPTIMIZE_CASE_EXPRESSION_PREDICATE, Boolean.class);
+    }
+
     public static boolean isDistributedSortEnabled(Session session)
     {
         return session.getSystemProperty(DISTRIBUTED_SORT, Boolean.class);

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java b/presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java
@@ -95,6 +95,7 @@ public class FeaturesConfig
     private int optimizeMetadataQueriesCallThreshold = 100;
     private boolean optimizeHashGeneration = true;
     private boolean enableIntermediateAggregations;
+    private boolean optimizeCaseExpressionPredicate;
     private boolean pushTableWriteThroughUnion = true;
     private boolean exchangeCompressionEnabled;
     private boolean exchangeChecksumEnabled;
@@ -874,6 +875,18 @@ public FeaturesConfig setOptimizeTopNRowNumber(boolean optimizeTopNRowNumber)
         return this;
     }
 
+    public boolean isOptimizeCaseExpressionPredicate()
+    {
+        return optimizeCaseExpressionPredicate;
+    }
+
+    @Config("optimizer.optimize-case-expression-predicate")
+    public FeaturesConfig setOptimizeCaseExpressionPredicate(boolean optimizeCaseExpressionPredicate)
+    {
+        this.optimizeCaseExpressionPredicate = optimizeCaseExpressionPredicate;
+        return this;
+    }
+
     public boolean isOptimizeHashGeneration()
     {
         return optimizeHashGeneration;

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java
@@ -98,6 +98,7 @@
 import com.facebook.presto.sql.planner.iterative.rule.RemoveUnsupportedDynamicFilters;
 import com.facebook.presto.sql.planner.iterative.rule.ReorderJoins;
 import com.facebook.presto.sql.planner.iterative.rule.RewriteAggregationIfToFilter;
+import com.facebook.presto.sql.planner.iterative.rule.RewriteCaseExpressionPredicate;
 import com.facebook.presto.sql.planner.iterative.rule.RewriteFilterWithExternalFunctionToProject;
 import com.facebook.presto.sql.planner.iterative.rule.RewriteSpatialPartitioningAggregation;
 import com.facebook.presto.sql.planner.iterative.rule.RuntimeReorderJoinSides;
@@ -277,6 +278,12 @@ public PlanOptimizers(
                         .add(new PruneRedundantProjectionAssignments())
                         .build());
 
+        IterativeOptimizer caseExpressionPredicateRewriter = new IterativeOptimizer(
+                ruleStats,
+                statsCalculator,
+                estimatedExchangesCostCalculator,
+                new RewriteCaseExpressionPredicate(metadata.getFunctionAndTypeManager()).rules());
+
         PlanOptimizer predicatePushDown = new StatsRecordingPlanOptimizer(optimizerStats, new PredicatePushDown(metadata, sqlParser));
 
         builder.add(
@@ -405,6 +412,7 @@ public PlanOptimizers(
         // After this point, all planNodes should not contain OriginalExpression
 
         builder.add(
+                caseExpressionPredicateRewriter,
                 new IterativeOptimizer(
                         ruleStats,
                         statsCalculator,