feat(optimizer): Add WindowFilterPushdown rule changes for rank queries

Author: aditi-pandit

presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -182,6 +182,7 @@ public final class SystemSessionProperties
     public static final String ADAPTIVE_PARTIAL_AGGREGATION = "adaptive_partial_aggregation";
     public static final String ADAPTIVE_PARTIAL_AGGREGATION_ROWS_REDUCTION_RATIO_THRESHOLD = "adaptive_partial_aggregation_unique_rows_ratio_threshold";
     public static final String OPTIMIZE_TOP_N_ROW_NUMBER = "optimize_top_n_row_number";
+    public static final String OPTIMIZE_TOP_N_RANK = "optimize_top_n_rank";
     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";
@@ -980,6 +981,11 @@ public SystemSessionProperties(
                         "Use top N row number optimization",
                         featuresConfig.isOptimizeTopNRowNumber(),
                         false),
+                booleanProperty(
+                        OPTIMIZE_TOP_N_RANK,
+                        "Use top N rank and dense_rank optimization",
+                        featuresConfig.isOptimizeTopNRank(),
+                        false),
                 booleanProperty(
                         OPTIMIZE_CASE_EXPRESSION_PREDICATE,
                         "Optimize case expression predicates",
@@ -2555,6 +2561,11 @@ public static boolean isOptimizeTopNRowNumber(Session session)
         return session.getSystemProperty(OPTIMIZE_TOP_N_ROW_NUMBER, Boolean.class);
     }
 
+    public static boolean isOptimizeTopNRank(Session session)
+    {
+        return session.getSystemProperty(OPTIMIZE_TOP_N_RANK, Boolean.class);
+    }
+
     public static boolean isOptimizeCaseExpressionPredicate(Session session)
     {
         return session.getSystemProperty(OPTIMIZE_CASE_EXPRESSION_PREDICATE, Boolean.class);

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java

@@ -163,6 +163,8 @@ public class FeaturesConfig
     private boolean adaptivePartialAggregationEnabled;
     private double adaptivePartialAggregationRowsReductionRatioThreshold = 0.8;
     private boolean optimizeTopNRowNumber = true;
+
+    private boolean optimizeTopNRank;
     private boolean pushLimitThroughOuterJoin = true;
     private boolean optimizeConstantGroupingKeys = true;
 
@@ -1141,13 +1143,25 @@ public boolean isOptimizeTopNRowNumber()
         return optimizeTopNRowNumber;
     }
 
+    public boolean isOptimizeTopNRank()
+    {
+        return optimizeTopNRank;
+    }
+
     @Config("optimizer.optimize-top-n-row-number")
     public FeaturesConfig setOptimizeTopNRowNumber(boolean optimizeTopNRowNumber)
     {
         this.optimizeTopNRowNumber = optimizeTopNRowNumber;
         return this;
     }
 
+    @Config("optimizer.optimize-top-n-rank")
+    public FeaturesConfig setOptimizeTopNRank(boolean optimizeTopNRank)
+    {
+        this.optimizeTopNRank = optimizeTopNRank;
+        return this;
+    }
+
     public boolean isOptimizeCaseExpressionPredicate()
     {
         return optimizeCaseExpressionPredicate;

presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java

@@ -41,11 +41,14 @@
 import com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator;
 import com.facebook.presto.sql.relational.RowExpressionDomainTranslator;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 
 import java.util.Map;
 import java.util.Optional;
 import java.util.OptionalInt;
 
+import static com.facebook.presto.SystemSessionProperties.isNativeExecutionEnabled;
+import static com.facebook.presto.SystemSessionProperties.isOptimizeTopNRank;
 import static com.facebook.presto.SystemSessionProperties.isOptimizeTopNRowNumber;
 import static com.facebook.presto.common.predicate.Marker.Bound.BELOW;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
@@ -134,6 +137,12 @@ public PlanNode visitWindow(WindowNode node, RewriteContext<Void> context)
             return replaceChildren(node, ImmutableList.of(rewrittenSource));
         }
 
+        private boolean canReplaceWithTopNRowNumber(WindowNode node)
+        {
+            return (canOptimizeRowNumberFunction(node, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) ||
+                    (canOptimizeRankFunction(node, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRank(session) && isNativeExecutionEnabled(session));
+        }
+
         @Override
         public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context)
         {
@@ -152,16 +161,16 @@ public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context)
                 planChanged = true;
                 source = rowNumberNode;
             }
-            else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
+            else if (source instanceof WindowNode) {
                 WindowNode windowNode = (WindowNode) source;
-                // verify that unordered row_number window functions are replaced by RowNumberNode
-                verify(windowNode.getOrderingScheme().isPresent());
-                TopNRowNumberNode topNRowNumberNode = convertToTopNRowNumber(windowNode, limit);
-                if (windowNode.getPartitionBy().isEmpty()) {
-                    return topNRowNumberNode;
+                if (canReplaceWithTopNRowNumber(windowNode)) {
+                    TopNRowNumberNode topNRowNumberNode = convertToTopNRowNumber(windowNode, limit);
+                    if (windowNode.getPartitionBy().isEmpty()) {
+                        return topNRowNumberNode;
+                    }
+                    planChanged = true;
+                    source = topNRowNumberNode;
                 }
-                planChanged = true;
-                source = topNRowNumberNode;
             }
             return replaceChildren(node, ImmutableList.of(source));
         }
@@ -183,15 +192,17 @@ public PlanNode visitFilter(FilterNode node, RewriteContext<Void> context)
                     return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
                 }
             }
-            else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
+            else if (source instanceof WindowNode) {
                 WindowNode windowNode = (WindowNode) source;
-                VariableReferenceExpression rowNumberVariable = getOnlyElement(windowNode.getCreatedVariable());
-                OptionalInt upperBound = extractUpperBound(tupleDomain, rowNumberVariable);
-
-                if (upperBound.isPresent()) {
-                    source = convertToTopNRowNumber(windowNode, upperBound.getAsInt());
-                    planChanged = true;
-                    return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
+                if (canReplaceWithTopNRowNumber(windowNode)) {
+                    VariableReferenceExpression rowNumberVariable = getOnlyElement(windowNode.getCreatedVariable());
+                    OptionalInt upperBound = extractUpperBound(tupleDomain, rowNumberVariable);
+
+                    if (upperBound.isPresent()) {
+                        source = convertToTopNRowNumber(windowNode, upperBound.getAsInt());
+                        planChanged = true;
+                        return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
+                    }
                 }
             }
             return replaceChildren(node, ImmutableList.of(source));
@@ -287,12 +298,44 @@ private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limi
                     Optional.empty());
         }
 
+        private TopNRowNumberNode convertToTopNRank(WindowNode windowNode, int limit)
+        {
+            String windowFunction = Iterables.getOnlyElement(windowNode.getWindowFunctions().values()).getFunctionCall().getFunctionHandle().getName();
+            String[] parts = windowFunction.split("\\.");
+            String windowFunctionName = parts[parts.length - 1];
+            TopNRowNumberNode.RankingFunction rankingFunction;
+            switch (windowFunctionName) {
+                case "row_number":
+                    rankingFunction = TopNRowNumberNode.RankingFunction.ROW_NUMBER;
+                    break;
+                case "rank":
+                    rankingFunction = TopNRowNumberNode.RankingFunction.RANK;
+                    break;
+                case "dense_rank":
+                    rankingFunction = TopNRowNumberNode.RankingFunction.DENSE_RANK;
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unsupported window function for TopNRowNumberNode: " + windowFunctionName);
+            }
+
+            return new TopNRowNumberNode(
+                    windowNode.getSourceLocation(),
+                    idAllocator.getNextId(),
+                    windowNode.getSource(),
+                    windowNode.getSpecification(),
+                    rankingFunction,
+                    getOnlyElement(windowNode.getCreatedVariable()),
+                    limit,
+                    false,
+                    Optional.empty());
+        }
+
         private static boolean canReplaceWithRowNumber(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
         {
-            return canOptimizeWindowFunction(node, functionAndTypeManager) && !node.getOrderingScheme().isPresent();
+            return canOptimizeRowNumberFunction(node, functionAndTypeManager) && !node.getOrderingScheme().isPresent();
         }
 
-        private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
+        private static boolean canOptimizeRowNumberFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
         {
             if (node.getWindowFunctions().size() != 1) {
                 return false;
@@ -301,10 +344,27 @@ private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTyp
             return isRowNumberMetadata(functionAndTypeManager, functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()));
         }
 
+        private static boolean canOptimizeRankFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
+        {
+            if (node.getWindowFunctions().size() != 1) {
+                return false;
+            }
+            VariableReferenceExpression rowNumberVariable = getOnlyElement(node.getWindowFunctions().keySet());
+            return isRankMetadata(functionAndTypeManager, functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()));
+        }
+
         private static boolean isRowNumberMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
         {
             FunctionHandle rowNumberFunction = functionAndTypeManager.lookupFunction("row_number", ImmutableList.of());
             return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rowNumberFunction));
         }
+
+        private static boolean isRankMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
+        {
+            FunctionHandle rankFunction = functionAndTypeManager.lookupFunction("rank", ImmutableList.of());
+            FunctionHandle denseRankFunction = functionAndTypeManager.lookupFunction("dense_rank", ImmutableList.of());
+            return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rankFunction)) ||
+                    functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(denseRankFunction));
+        }
     }
 }

presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TopNRowNumberNode.java

@@ -85,6 +85,7 @@ public TopNRowNumberNode(
         requireNonNull(rowNumberVariable, "rowNumberVariable is null");
         checkArgument(maxRowCountPerPartition > 0, "maxRowCountPerPartition must be > 0");
         requireNonNull(hashVariable, "hashVariable is null");
+        requireNonNull(rankingFunction, "rankingFunction is null");
 
         this.source = source;
         this.specification = specification;

presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java

@@ -145,6 +145,7 @@ public void testDefaults()
                 .setAdaptivePartialAggregationEnabled(false)
                 .setAdaptivePartialAggregationRowsReductionRatioThreshold(0.8)
                 .setOptimizeTopNRowNumber(true)
+                .setOptimizeTopNRank(false)
                 .setOptimizeCaseExpressionPredicate(false)
                 .setDistributedSortEnabled(true)
                 .setMaxGroupingSets(2048)
@@ -367,6 +368,7 @@ public void testExplicitPropertyMappings()
                 .put("experimental.adaptive-partial-aggregation", "true")
                 .put("experimental.adaptive-partial-aggregation-rows-reduction-ratio-threshold", "0.9")
                 .put("optimizer.optimize-top-n-row-number", "false")
+                .put("optimizer.optimize-top-n-rank", "true")
                 .put("optimizer.optimize-case-expression-predicate", "true")
                 .put("distributed-sort", "false")
                 .put("analyzer.max-grouping-sets", "2047")
@@ -584,6 +586,7 @@ public void testExplicitPropertyMappings()
                 .setAdaptivePartialAggregationEnabled(true)
                 .setAdaptivePartialAggregationRowsReductionRatioThreshold(0.9)
                 .setOptimizeTopNRowNumber(false)
+                .setOptimizeTopNRank(true)
                 .setOptimizeCaseExpressionPredicate(true)
                 .setDistributedSortEnabled(false)
                 .setMaxGroupingSets(2047)