Skip to content

Commit 5013a95

Browse files
aditi-panditaditi-pandit
committed
feat(optimizer): Add WindowFilterPushdown rule changes for rank queries
1 parent 216dadc commit 5013a95

File tree

7 files changed

+227
-20
lines changed

7 files changed

+227
-20
lines changed

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,7 @@ public final class SystemSessionProperties
182182
public static final String ADAPTIVE_PARTIAL_AGGREGATION = "adaptive_partial_aggregation";
183183
public static final String ADAPTIVE_PARTIAL_AGGREGATION_ROWS_REDUCTION_RATIO_THRESHOLD = "adaptive_partial_aggregation_unique_rows_ratio_threshold";
184184
public static final String OPTIMIZE_TOP_N_ROW_NUMBER = "optimize_top_n_row_number";
185+
public static final String OPTIMIZE_TOP_N_RANK = "optimize_top_n_rank";
185186
public static final String OPTIMIZE_CASE_EXPRESSION_PREDICATE = "optimize_case_expression_predicate";
186187
public static final String MAX_GROUPING_SETS = "max_grouping_sets";
187188
public static final String LEGACY_UNNEST = "legacy_unnest";
@@ -980,6 +981,11 @@ public SystemSessionProperties(
980981
"Use top N row number optimization",
981982
featuresConfig.isOptimizeTopNRowNumber(),
982983
false),
984+
booleanProperty(
985+
OPTIMIZE_TOP_N_RANK,
986+
"Use top N rank and dense_rank optimization",
987+
featuresConfig.isOptimizeTopNRank(),
988+
false),
983989
booleanProperty(
984990
OPTIMIZE_CASE_EXPRESSION_PREDICATE,
985991
"Optimize case expression predicates",
@@ -2555,6 +2561,11 @@ public static boolean isOptimizeTopNRowNumber(Session session)
25552561
return session.getSystemProperty(OPTIMIZE_TOP_N_ROW_NUMBER, Boolean.class);
25562562
}
25572563

2564+
public static boolean isOptimizeTopNRank(Session session)
2565+
{
2566+
return session.getSystemProperty(OPTIMIZE_TOP_N_RANK, Boolean.class);
2567+
}
2568+
25582569
public static boolean isOptimizeCaseExpressionPredicate(Session session)
25592570
{
25602571
return session.getSystemProperty(OPTIMIZE_CASE_EXPRESSION_PREDICATE, Boolean.class);

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,8 @@ public class FeaturesConfig
163163
private boolean adaptivePartialAggregationEnabled;
164164
private double adaptivePartialAggregationRowsReductionRatioThreshold = 0.8;
165165
private boolean optimizeTopNRowNumber = true;
166+
167+
private boolean optimizeTopNRank;
166168
private boolean pushLimitThroughOuterJoin = true;
167169
private boolean optimizeConstantGroupingKeys = true;
168170

@@ -1135,19 +1137,30 @@ public FeaturesConfig setAdaptivePartialAggregationRowsReductionRatioThreshold(d
11351137
this.adaptivePartialAggregationRowsReductionRatioThreshold = adaptivePartialAggregationRowsReductionRatioThreshold;
11361138
return this;
11371139
}
1138-
11391140
public boolean isOptimizeTopNRowNumber()
11401141
{
11411142
return optimizeTopNRowNumber;
11421143
}
11431144

1145+
public boolean isOptimizeTopNRank()
1146+
{
1147+
return optimizeTopNRank;
1148+
}
1149+
11441150
@Config("optimizer.optimize-top-n-row-number")
11451151
public FeaturesConfig setOptimizeTopNRowNumber(boolean optimizeTopNRowNumber)
11461152
{
11471153
this.optimizeTopNRowNumber = optimizeTopNRowNumber;
11481154
return this;
11491155
}
11501156

1157+
@Config("optimizer.optimize-top-n-rank")
1158+
public FeaturesConfig setOptimizeTopNRank(boolean optimizeTopNRank)
1159+
{
1160+
this.optimizeTopNRank = optimizeTopNRank;
1161+
return this;
1162+
}
1163+
11511164
public boolean isOptimizeCaseExpressionPredicate()
11521165
{
11531166
return optimizeCaseExpressionPredicate;

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

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,14 @@
4141
import com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator;
4242
import com.facebook.presto.sql.relational.RowExpressionDomainTranslator;
4343
import com.google.common.collect.ImmutableList;
44+
import com.google.common.collect.Iterables;
4445

4546
import java.util.Map;
4647
import java.util.Optional;
4748
import java.util.OptionalInt;
4849

50+
import static com.facebook.presto.SystemSessionProperties.isNativeExecutionEnabled;
51+
import static com.facebook.presto.SystemSessionProperties.isOptimizeTopNRank;
4952
import static com.facebook.presto.SystemSessionProperties.isOptimizeTopNRowNumber;
5053
import static com.facebook.presto.common.predicate.Marker.Bound.BELOW;
5154
import static com.facebook.presto.common.type.BigintType.BIGINT;
@@ -152,11 +155,12 @@ public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context)
152155
planChanged = true;
153156
source = rowNumberNode;
154157
}
155-
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
158+
else if ((source instanceof WindowNode && canOptimizeRowNumberFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) ||
159+
(source instanceof WindowNode && canOptimizeRankFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRank(session) && isNativeExecutionEnabled(session))) {
156160
WindowNode windowNode = (WindowNode) source;
157161
// verify that unordered row_number window functions are replaced by RowNumberNode
158162
verify(windowNode.getOrderingScheme().isPresent());
159-
TopNRowNumberNode topNRowNumberNode = convertToTopNRowNumber(windowNode, limit);
163+
TopNRowNumberNode topNRowNumberNode = convertToTopNRank(windowNode, limit);
160164
if (windowNode.getPartitionBy().isEmpty()) {
161165
return topNRowNumberNode;
162166
}
@@ -183,7 +187,8 @@ public PlanNode visitFilter(FilterNode node, RewriteContext<Void> context)
183187
return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
184188
}
185189
}
186-
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
190+
else if ((source instanceof WindowNode && canOptimizeRowNumberFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) ||
191+
(source instanceof WindowNode && canOptimizeRankFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRank(session) && isNativeExecutionEnabled(session))) {
187192
WindowNode windowNode = (WindowNode) source;
188193
VariableReferenceExpression rowNumberVariable = getOnlyElement(windowNode.getCreatedVariable());
189194
OptionalInt upperBound = extractUpperBound(tupleDomain, rowNumberVariable);
@@ -280,19 +285,52 @@ private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limi
280285
idAllocator.getNextId(),
281286
windowNode.getSource(),
282287
windowNode.getSpecification(),
288+
283289
TopNRowNumberNode.RankingFunction.ROW_NUMBER,
284290
getOnlyElement(windowNode.getCreatedVariable()),
285291
limit,
286292
false,
287293
Optional.empty());
288294
}
289295

296+
private TopNRowNumberNode convertToTopNRank(WindowNode windowNode, int limit)
297+
{
298+
String windowFunction = Iterables.getOnlyElement(windowNode.getWindowFunctions().values()).getFunctionCall().getFunctionHandle().getName();
299+
String[] parts = windowFunction.split("\\.");
300+
String windowFunctionName = parts[parts.length - 1];
301+
TopNRowNumberNode.RankingFunction rankingFunction;
302+
switch (windowFunctionName) {
303+
case "row_number":
304+
rankingFunction = TopNRowNumberNode.RankingFunction.ROW_NUMBER;
305+
break;
306+
case "rank":
307+
rankingFunction = TopNRowNumberNode.RankingFunction.RANK;
308+
break;
309+
case "dense_rank":
310+
rankingFunction = TopNRowNumberNode.RankingFunction.DENSE_RANK;
311+
break;
312+
default:
313+
throw new IllegalArgumentException("Unsupported window function for TopNRowNumberNode: " + windowFunctionName);
314+
}
315+
316+
return new TopNRowNumberNode(
317+
windowNode.getSourceLocation(),
318+
idAllocator.getNextId(),
319+
windowNode.getSource(),
320+
windowNode.getSpecification(),
321+
rankingFunction,
322+
getOnlyElement(windowNode.getCreatedVariable()),
323+
limit,
324+
false,
325+
Optional.empty());
326+
}
327+
290328
private static boolean canReplaceWithRowNumber(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
291329
{
292-
return canOptimizeWindowFunction(node, functionAndTypeManager) && !node.getOrderingScheme().isPresent();
330+
return canOptimizeRowNumberFunction(node, functionAndTypeManager) && !node.getOrderingScheme().isPresent();
293331
}
294332

295-
private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
333+
private static boolean canOptimizeRowNumberFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
296334
{
297335
if (node.getWindowFunctions().size() != 1) {
298336
return false;
@@ -301,10 +339,27 @@ private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTyp
301339
return isRowNumberMetadata(functionAndTypeManager, functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()));
302340
}
303341

342+
private static boolean canOptimizeRankFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
343+
{
344+
if (node.getWindowFunctions().size() != 1) {
345+
return false;
346+
}
347+
VariableReferenceExpression rowNumberVariable = getOnlyElement(node.getWindowFunctions().keySet());
348+
return isRankMetadata(functionAndTypeManager, functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()));
349+
}
350+
304351
private static boolean isRowNumberMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
305352
{
306353
FunctionHandle rowNumberFunction = functionAndTypeManager.lookupFunction("row_number", ImmutableList.of());
307354
return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rowNumberFunction));
308355
}
356+
357+
private static boolean isRankMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
358+
{
359+
FunctionHandle rankFunction = functionAndTypeManager.lookupFunction("rank", ImmutableList.of());
360+
FunctionHandle denseRankFunction = functionAndTypeManager.lookupFunction("dense_rank", ImmutableList.of());
361+
return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rankFunction)) ||
362+
functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(denseRankFunction));
363+
}
309364
}
310365
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@ public void testDefaults()
145145
.setAdaptivePartialAggregationEnabled(false)
146146
.setAdaptivePartialAggregationRowsReductionRatioThreshold(0.8)
147147
.setOptimizeTopNRowNumber(true)
148+
.setOptimizeTopNRank(false)
148149
.setOptimizeCaseExpressionPredicate(false)
149150
.setDistributedSortEnabled(true)
150151
.setMaxGroupingSets(2048)
@@ -367,6 +368,7 @@ public void testExplicitPropertyMappings()
367368
.put("experimental.adaptive-partial-aggregation", "true")
368369
.put("experimental.adaptive-partial-aggregation-rows-reduction-ratio-threshold", "0.9")
369370
.put("optimizer.optimize-top-n-row-number", "false")
371+
.put("optimizer.optimize-top-n-rank", "true")
370372
.put("optimizer.optimize-case-expression-predicate", "true")
371373
.put("distributed-sort", "false")
372374
.put("analyzer.max-grouping-sets", "2047")
@@ -584,6 +586,7 @@ public void testExplicitPropertyMappings()
584586
.setAdaptivePartialAggregationEnabled(true)
585587
.setAdaptivePartialAggregationRowsReductionRatioThreshold(0.9)
586588
.setOptimizeTopNRowNumber(false)
589+
.setOptimizeTopNRank(true)
587590
.setOptimizeCaseExpressionPredicate(true)
588591
.setDistributedSortEnabled(false)
589592
.setMaxGroupingSets(2047)

‎presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestWindowFilterPushDown.java‎

Lines changed: 82 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import org.intellij.lang.annotations.Language;
2222
import org.testng.annotations.Test;
2323

24+
import static com.facebook.presto.SystemSessionProperties.NATIVE_EXECUTION_ENABLED;
25+
import static com.facebook.presto.SystemSessionProperties.OPTIMIZE_TOP_N_RANK;
2426
import static com.facebook.presto.SystemSessionProperties.OPTIMIZE_TOP_N_ROW_NUMBER;
2527
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyNot;
2628
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree;
@@ -31,15 +33,11 @@
3133
public class TestWindowFilterPushDown
3234
extends BasePlanTest
3335
{
34-
@Test
35-
public void testLimitAboveWindow()
36+
private void testLimitSql(String sql, boolean rowOrRank)
3637
{
37-
@Language("SQL") String sql = "SELECT " +
38-
"row_number() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem LIMIT 10";
39-
4038
assertPlanWithSession(
4139
sql,
42-
optimizeTopNRowNumber(true),
40+
rowOrRank ? optimizeTopNRowNumber(true) : optimizeTopNRank(true),
4341
true,
4442
anyTree(
4543
limit(10, anyTree(
@@ -49,25 +47,47 @@ public void testLimitAboveWindow()
4947

5048
assertPlanWithSession(
5149
sql,
52-
optimizeTopNRowNumber(false),
50+
rowOrRank ? optimizeTopNRowNumber(false) : optimizeTopNRank(false),
5351
true,
5452
anyTree(
5553
limit(10, anyTree(
5654
node(WindowNode.class,
5755
anyTree(
5856
tableScan("lineitem")))))));
59-
}
6057

58+
if (!rowOrRank) {
59+
assertPlanWithSession(
60+
sql,
61+
optimizeTopNRankWithoutNative(true),
62+
true,
63+
anyTree(
64+
limit(10, anyTree(
65+
node(WindowNode.class,
66+
anyTree(
67+
tableScan("lineitem")))))));
68+
}
69+
}
6170
@Test
62-
public void testFilterAboveWindow()
71+
public void testLimitAboveWindow()
6372
{
64-
@Language("SQL") String sql = "SELECT * FROM " +
65-
"(SELECT row_number() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem) " +
66-
"WHERE partition_row_number < 10";
73+
@Language("SQL") String sql = "SELECT " +
74+
"row_number() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem LIMIT 10";
75+
testLimitSql(sql, true);
76+
77+
sql = "SELECT " +
78+
"rank() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem LIMIT 10";
79+
testLimitSql(sql, false);
6780

81+
sql = "SELECT " +
82+
"dense_rank() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem LIMIT 10";
83+
testLimitSql(sql, false);
84+
}
85+
86+
private void testFilterSql(String sql, boolean rowOrRank)
87+
{
6888
assertPlanWithSession(
6989
sql,
70-
optimizeTopNRowNumber(true),
90+
rowOrRank ? optimizeTopNRowNumber(true) : optimizeTopNRank(true),
7191
true,
7292
anyTree(
7393
anyNot(FilterNode.class,
@@ -77,14 +97,46 @@ public void testFilterAboveWindow()
7797

7898
assertPlanWithSession(
7999
sql,
80-
optimizeTopNRowNumber(false),
100+
rowOrRank ? optimizeTopNRowNumber(false) : optimizeTopNRank(false),
81101
true,
82102
anyTree(
83103
node(FilterNode.class,
84104
anyTree(
85105
node(WindowNode.class,
86106
anyTree(
87107
tableScan("lineitem")))))));
108+
109+
if (!rowOrRank) {
110+
assertPlanWithSession(
111+
sql,
112+
optimizeTopNRankWithoutNative(true),
113+
true,
114+
anyTree(
115+
node(FilterNode.class,
116+
anyTree(
117+
node(WindowNode.class,
118+
anyTree(
119+
tableScan("lineitem")))))));
120+
}
121+
}
122+
@Test
123+
public void testFilterAboveWindow()
124+
{
125+
@Language("SQL") String sql = "SELECT * FROM " +
126+
"(SELECT row_number() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_row_number FROM lineitem) " +
127+
"WHERE partition_row_number < 10";
128+
129+
testFilterSql(sql, true);
130+
131+
sql = "SELECT * FROM " +
132+
"(SELECT rank() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_rank FROM lineitem) " +
133+
"WHERE partition_rank < 10";
134+
testFilterSql(sql, false);
135+
136+
sql = "SELECT * FROM " +
137+
"(SELECT dense_rank() OVER (PARTITION BY suppkey ORDER BY orderkey) partition_dense_rank FROM lineitem) " +
138+
"WHERE partition_dense_rank < 10";
139+
testFilterSql(sql, false);
88140
}
89141

90142
private Session optimizeTopNRowNumber(boolean enabled)
@@ -93,4 +145,20 @@ private Session optimizeTopNRowNumber(boolean enabled)
93145
.setSystemProperty(OPTIMIZE_TOP_N_ROW_NUMBER, Boolean.toString(enabled))
94146
.build();
95147
}
148+
149+
private Session optimizeTopNRank(boolean enabled)
150+
{
151+
return Session.builder(this.getQueryRunner().getDefaultSession())
152+
.setSystemProperty(NATIVE_EXECUTION_ENABLED, Boolean.toString(enabled))
153+
.setSystemProperty(OPTIMIZE_TOP_N_RANK, Boolean.toString(enabled))
154+
.build();
155+
}
156+
157+
private Session optimizeTopNRankWithoutNative(boolean enabled)
158+
{
159+
return Session.builder(this.getQueryRunner().getDefaultSession())
160+
.setSystemProperty(NATIVE_EXECUTION_ENABLED, Boolean.toString(false))
161+
.setSystemProperty(OPTIMIZE_TOP_N_RANK, Boolean.toString(enabled))
162+
.build();
163+
}
96164
}

0 commit comments

Comments
 (0)