Skip to content

Commit 99600e4

Browse files
aditi-panditaditi-pandit
committed
[native] TopNRank optimization
1 parent 461ae13 commit 99600e4

File tree

16 files changed

+226
-24
lines changed

16 files changed

+226
-24
lines changed

‎presto-main/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -696,6 +696,7 @@ public Optional<PlanNode> visitTopNRowNumber(TopNRowNumberNode node, Context con
696696
new WindowNode.Specification(
697697
partitionBy,
698698
node.getSpecification().getOrderingScheme().map(scheme -> getCanonicalOrderingScheme(scheme, context.getExpressions()))),
699+
node.getRankingFunction(),
699700
rowNumberVariable,
700701
node.getMaxRowCountPerPartition(),
701702
node.isPartial(),

‎presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -641,7 +641,7 @@ public PlanOptimizers(
641641
estimatedExchangesCostCalculator,
642642
ImmutableSet.of(new SimplifyCountOverConstant(metadata.getFunctionAndTypeManager()))),
643643
new LimitPushDown(), // Run LimitPushDown before WindowFilterPushDown
644-
new WindowFilterPushDown(metadata), // This must run after PredicatePushDown and LimitPushDown so that it squashes any successive filter nodes and limits
644+
new WindowFilterPushDown(metadata, featuresConfig.isNativeExecutionEnabled()), // This must run after PredicatePushDown and LimitPushDown so that it squashes any successive filter nodes and limits
645645
prefilterForLimitingAggregation,
646646
new IterativeOptimizer(
647647
metadata,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,7 @@ public PlanWithProperties visitTopNRowNumber(TopNRowNumberNode node, PreferredPr
493493
idAllocator.getNextId(),
494494
child.getNode(),
495495
node.getSpecification(),
496+
node.getRankingFunction(),
496497
node.getRowNumberVariable(),
497498
node.getMaxRowCountPerPartition(),
498499
true,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,7 @@ public PlanWithProperties visitTopNRowNumber(TopNRowNumberNode node, HashComputa
329329
node.getId(),
330330
child.getNode(),
331331
node.getSpecification(),
332+
node.getRankingFunction(),
332333
node.getRowNumberVariable(),
333334
node.getMaxRowCountPerPartition(),
334335
node.isPartial(),

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -310,6 +310,7 @@ public Optional<DecorrelationResult> visitTopN(TopNNode node, Void context)
310310
new Specification(
311311
ImmutableList.copyOf(childDecorrelationResult.variablesToPropagate),
312312
Optional.of(orderingScheme)),
313+
TopNRowNumberNode.RankingFunction.ROW_NUMBER,
313314
variableAllocator.newVariable("row_number", BIGINT),
314315
toIntExact(node.getCount()),
315316
false,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -701,6 +701,7 @@ public PlanNode visitTopNRowNumber(TopNRowNumberNode node, RewriteContext<Set<Va
701701
node.getStatsEquivalentPlanNode(),
702702
source,
703703
node.getSpecification(),
704+
node.getRankingFunction(),
704705
node.getRowNumberVariable(),
705706
node.getMaxRowCountPerPartition(),
706707
node.isPartial(),

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,7 @@ public PlanNode visitTopNRowNumber(TopNRowNumberNode node, RewriteContext<Void>
481481
node.getId(),
482482
context.rewrite(node.getSource()),
483483
canonicalizeAndDistinct(node.getSpecification()),
484+
node.getRankingFunction(),
484485
canonicalize(node.getRowNumberVariable()),
485486
node.getMaxRowCountPerPartition(),
486487
node.isPartial(),

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

Lines changed: 55 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import java.util.Optional;
4747
import java.util.OptionalInt;
4848

49+
import static com.facebook.presto.SystemSessionProperties.isNativeExecutionEnabled;
4950
import static com.facebook.presto.SystemSessionProperties.isOptimizeTopNRowNumber;
5051
import static com.facebook.presto.common.predicate.Marker.Bound.BELOW;
5152
import static com.facebook.presto.common.type.BigintType.BIGINT;
@@ -65,10 +66,13 @@ public class WindowFilterPushDown
6566
private final RowExpressionDomainTranslator domainTranslator;
6667
private final LogicalRowExpressions logicalRowExpressions;
6768

68-
public WindowFilterPushDown(Metadata metadata)
69+
private boolean isNativeExecution = false;
70+
71+
public WindowFilterPushDown(Metadata metadata, boolean isNativeExecution)
6972
{
7073
this.metadata = requireNonNull(metadata, "metadata is null");
7174
this.domainTranslator = new RowExpressionDomainTranslator(metadata);
75+
this.isNativeExecution = isNativeExecution;
7276
this.logicalRowExpressions = new LogicalRowExpressions(
7377
new RowExpressionDeterminismEvaluator(metadata.getFunctionAndTypeManager()),
7478
new FunctionResolution(metadata.getFunctionAndTypeManager().getFunctionAndTypeResolver()),
@@ -84,7 +88,7 @@ public PlanOptimizerResult optimize(PlanNode plan, Session session, TypeProvider
8488
requireNonNull(variableAllocator, "variableAllocator is null");
8589
requireNonNull(idAllocator, "idAllocator is null");
8690

87-
Rewriter rewriter = new Rewriter(idAllocator, metadata, domainTranslator, logicalRowExpressions, session);
91+
Rewriter rewriter = new Rewriter(idAllocator, metadata, domainTranslator, logicalRowExpressions, session, isNativeExecution);
8892
PlanNode rewrittenPlan = SimplePlanRewriter.rewriteWith(rewriter, plan, null);
8993
return PlanOptimizerResult.optimizerResult(rewrittenPlan, rewriter.isPlanChanged());
9094
}
@@ -97,15 +101,17 @@ private static class Rewriter
97101
private final RowExpressionDomainTranslator domainTranslator;
98102
private final LogicalRowExpressions logicalRowExpressions;
99103
private final Session session;
104+
private final boolean isNativeExecution;
100105
private boolean planChanged;
101106

102-
private Rewriter(PlanNodeIdAllocator idAllocator, Metadata metadata, RowExpressionDomainTranslator domainTranslator, LogicalRowExpressions logicalRowExpressions, Session session)
107+
private Rewriter(PlanNodeIdAllocator idAllocator, Metadata metadata, RowExpressionDomainTranslator domainTranslator, LogicalRowExpressions logicalRowExpressions, Session session, boolean isNativeExecution)
103108
{
104109
this.idAllocator = requireNonNull(idAllocator, "idAllocator is null");
105110
this.metadata = requireNonNull(metadata, "metadata is null");
106111
this.domainTranslator = requireNonNull(domainTranslator, "domainTranslator is null");
107112
this.logicalRowExpressions = logicalRowExpressions;
108113
this.session = requireNonNull(session, "session is null");
114+
this.isNativeExecution = isNativeExecution;
109115
}
110116

111117
public boolean isPlanChanged()
@@ -138,6 +144,7 @@ public PlanNode visitWindow(WindowNode node, RewriteContext<Void> context)
138144
public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context)
139145
{
140146
// Operators can handle MAX_VALUE rows per page, so do not optimize if count is greater than this value
147+
// TODO (Aditi) : Don't think this check is needed for Native engine.
141148
if (node.getCount() > Integer.MAX_VALUE) {
142149
return context.defaultRewrite(node);
143150
}
@@ -152,11 +159,11 @@ public PlanNode visitLimit(LimitNode node, RewriteContext<Void> context)
152159
planChanged = true;
153160
source = rowNumberNode;
154161
}
155-
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
162+
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager(), isNativeExecution) && isOptimizeTopNRowNumber(session)) {
156163
WindowNode windowNode = (WindowNode) source;
157164
// verify that unordered row_number window functions are replaced by RowNumberNode
158165
verify(windowNode.getOrderingScheme().isPresent());
159-
TopNRowNumberNode topNRowNumberNode = convertToTopNRowNumber(windowNode, limit);
166+
TopNRowNumberNode topNRowNumberNode = convertToTopNRowNumber(windowNode, limit, metadata.getFunctionAndTypeManager());
160167
if (windowNode.getPartitionBy().isEmpty()) {
161168
return topNRowNumberNode;
162169
}
@@ -183,13 +190,13 @@ public PlanNode visitFilter(FilterNode node, RewriteContext<Void> context)
183190
return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
184191
}
185192
}
186-
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager()) && isOptimizeTopNRowNumber(session)) {
193+
else if (source instanceof WindowNode && canOptimizeWindowFunction((WindowNode) source, metadata.getFunctionAndTypeManager(), isNativeExecution) && isOptimizeTopNRowNumber(session)) {
187194
WindowNode windowNode = (WindowNode) source;
188195
VariableReferenceExpression rowNumberVariable = getOnlyElement(windowNode.getCreatedVariable());
189196
OptionalInt upperBound = extractUpperBound(tupleDomain, rowNumberVariable);
190197

191198
if (upperBound.isPresent()) {
192-
source = convertToTopNRowNumber(windowNode, upperBound.getAsInt());
199+
source = convertToTopNRowNumber(windowNode, upperBound.getAsInt(), metadata.getFunctionAndTypeManager());
193200
planChanged = true;
194201
return rewriteFilterSource(node, source, rowNumberVariable, upperBound.getAsInt());
195202
}
@@ -273,13 +280,23 @@ private static RowNumberNode mergeLimit(RowNumberNode node, int newRowCountPerPa
273280
return new RowNumberNode(node.getSourceLocation(), node.getId(), node.getSource(), node.getPartitionBy(), node.getRowNumberVariable(), Optional.of(newRowCountPerPartition), false, node.getHashVariable());
274281
}
275282

276-
private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limit)
283+
private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limit, FunctionAndTypeManager functionAndTypeManager)
277284
{
285+
VariableReferenceExpression rowNumberVariable = getOnlyElement(windowNode.getWindowFunctions().keySet());
286+
FunctionMetadata functionMetadata = functionAndTypeManager.getFunctionMetadata(windowNode.getWindowFunctions().get(rowNumberVariable).getFunctionHandle());
287+
288+
TopNRowNumberNode.RankingFunction rankingFunction =
289+
isRowNumberMetadata(functionAndTypeManager, functionMetadata) ?
290+
TopNRowNumberNode.RankingFunction.ROW_NUMBER :
291+
isRankMetadata(functionAndTypeManager, functionMetadata) ?
292+
TopNRowNumberNode.RankingFunction.RANK :
293+
TopNRowNumberNode.RankingFunction.DENSE_RANK;
278294
return new TopNRowNumberNode(
279295
windowNode.getSourceLocation(),
280296
idAllocator.getNextId(),
281297
windowNode.getSource(),
282298
windowNode.getSpecification(),
299+
rankingFunction,
283300
getOnlyElement(windowNode.getCreatedVariable()),
284301
limit,
285302
false,
@@ -288,22 +305,49 @@ private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limi
288305

289306
private static boolean canReplaceWithRowNumber(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
290307
{
291-
return canOptimizeWindowFunction(node, functionAndTypeManager) && !node.getOrderingScheme().isPresent();
308+
if (node.getWindowFunctions().size() != 1) {
309+
return false;
310+
}
311+
VariableReferenceExpression rowNumberVariable = getOnlyElement(node.getWindowFunctions().keySet());
312+
313+
return isRowNumberMetadata(functionAndTypeManager,
314+
functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()))
315+
&& !node.getOrderingScheme().isPresent();
292316
}
293317

294-
private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager)
318+
private static boolean canOptimizeWindowFunction(WindowNode node, FunctionAndTypeManager functionAndTypeManager, boolean isNativeExecution)
295319
{
296320
if (node.getWindowFunctions().size() != 1) {
297321
return false;
298322
}
299323
VariableReferenceExpression rowNumberVariable = getOnlyElement(node.getWindowFunctions().keySet());
300-
return isRowNumberMetadata(functionAndTypeManager, functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle()));
324+
FunctionMetadata functionMetadata = functionAndTypeManager.getFunctionMetadata(node.getWindowFunctions().get(rowNumberVariable).getFunctionHandle());
325+
if (isNativeExecution) {
326+
return isRowNumberMetadata(functionAndTypeManager, functionMetadata)
327+
|| node.getOrderingScheme().isPresent() && (isRankMetadata(functionAndTypeManager, functionMetadata)
328+
|| isDenseRankMetadata(functionAndTypeManager, functionMetadata));
329+
}
330+
331+
return isRowNumberMetadata(functionAndTypeManager, functionMetadata);
301332
}
302333

334+
303335
private static boolean isRowNumberMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
304336
{
305337
FunctionHandle rowNumberFunction = functionAndTypeManager.lookupFunction("row_number", ImmutableList.of());
306338
return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rowNumberFunction));
307339
}
340+
341+
private static boolean isRankMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
342+
{
343+
FunctionHandle rankFunction = functionAndTypeManager.lookupFunction("rank", ImmutableList.of());
344+
return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rankFunction));
345+
}
346+
347+
private static boolean isDenseRankMetadata(FunctionAndTypeManager functionAndTypeManager, FunctionMetadata functionMetadata)
348+
{
349+
FunctionHandle rankFunction = functionAndTypeManager.lookupFunction("dense_rank", ImmutableList.of());
350+
return functionMetadata.equals(functionAndTypeManager.getFunctionMetadata(rankFunction));
351+
}
308352
}
309353
}

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

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,18 @@
3636
public final class TopNRowNumberNode
3737
extends InternalPlanNode
3838
{
39+
public enum RankingFunction
40+
{
41+
ROW_NUMBER,
42+
RANK,
43+
DENSE_RANK
44+
}
45+
3946
private final PlanNode source;
4047
private final Specification specification;
48+
private final RankingFunction rankingFunction;
4149
private final VariableReferenceExpression rowNumberVariable;
42-
private final int maxRowCountPerPartition;
50+
private final int maxRankPerPartition;
4351
private final boolean partial;
4452
private final Optional<VariableReferenceExpression> hashVariable;
4553

@@ -49,12 +57,13 @@ public TopNRowNumberNode(
4957
@JsonProperty("id") PlanNodeId id,
5058
@JsonProperty("source") PlanNode source,
5159
@JsonProperty("specification") Specification specification,
60+
@JsonProperty("rankingType") RankingFunction rankingFunction,
5261
@JsonProperty("rowNumberVariable") VariableReferenceExpression rowNumberVariable,
53-
@JsonProperty("maxRowCountPerPartition") int maxRowCountPerPartition,
62+
@JsonProperty("maxRowCountPerPartition") int maxRankPerPartition,
5463
@JsonProperty("partial") boolean partial,
5564
@JsonProperty("hashVariable") Optional<VariableReferenceExpression> hashVariable)
5665
{
57-
this(sourceLocation, id, Optional.empty(), source, specification, rowNumberVariable, maxRowCountPerPartition, partial, hashVariable);
66+
this(sourceLocation, id, Optional.empty(), source, specification, rankingFunction, rowNumberVariable, maxRankPerPartition, partial, hashVariable);
5867
}
5968

6069
public TopNRowNumberNode(
@@ -63,8 +72,9 @@ public TopNRowNumberNode(
6372
Optional<PlanNode> statsEquivalentPlanNode,
6473
PlanNode source,
6574
Specification specification,
75+
RankingFunction rankingFunction,
6676
VariableReferenceExpression rowNumberVariable,
67-
int maxRowCountPerPartition,
77+
int maxRankPerPartition,
6878
boolean partial,
6979
Optional<VariableReferenceExpression> hashVariable)
7080
{
@@ -74,13 +84,14 @@ public TopNRowNumberNode(
7484
requireNonNull(specification, "specification is null");
7585
checkArgument(specification.getOrderingScheme().isPresent(), "specification orderingScheme is absent");
7686
requireNonNull(rowNumberVariable, "rowNumberVariable is null");
77-
checkArgument(maxRowCountPerPartition > 0, "maxRowCountPerPartition must be > 0");
87+
checkArgument(maxRankPerPartition > 0, "maxRowCountPerPartition must be > 0");
7888
requireNonNull(hashVariable, "hashVariable is null");
7989

8090
this.source = source;
8191
this.specification = specification;
92+
this.rankingFunction = rankingFunction;
8293
this.rowNumberVariable = rowNumberVariable;
83-
this.maxRowCountPerPartition = maxRowCountPerPartition;
94+
this.maxRankPerPartition = maxRankPerPartition;
8495
this.partial = partial;
8596
this.hashVariable = hashVariable;
8697
}
@@ -108,6 +119,12 @@ public PlanNode getSource()
108119
return source;
109120
}
110121

122+
@JsonProperty
123+
public RankingFunction getRankingFunction()
124+
{
125+
return rankingFunction;
126+
}
127+
111128
@JsonProperty
112129
public Specification getSpecification()
113130
{
@@ -133,7 +150,7 @@ public VariableReferenceExpression getRowNumberVariable()
133150
@JsonProperty
134151
public int getMaxRowCountPerPartition()
135152
{
136-
return maxRowCountPerPartition;
153+
return maxRankPerPartition;
137154
}
138155

139156
@JsonProperty
@@ -157,12 +174,12 @@ public <R, C> R accept(InternalPlanVisitor<R, C> visitor, C context)
157174
@Override
158175
public PlanNode replaceChildren(List<PlanNode> newChildren)
159176
{
160-
return new TopNRowNumberNode(getSourceLocation(), getId(), getStatsEquivalentPlanNode(), Iterables.getOnlyElement(newChildren), specification, rowNumberVariable, maxRowCountPerPartition, partial, hashVariable);
177+
return new TopNRowNumberNode(getSourceLocation(), getId(), getStatsEquivalentPlanNode(), Iterables.getOnlyElement(newChildren), specification, rankingFunction, rowNumberVariable, maxRankPerPartition, partial, hashVariable);
161178
}
162179

163180
@Override
164181
public PlanNode assignStatsEquivalentPlanNode(Optional<PlanNode> statsEquivalentPlanNode)
165182
{
166-
return new TopNRowNumberNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, specification, rowNumberVariable, maxRowCountPerPartition, partial, hashVariable);
183+
return new TopNRowNumberNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, specification, rankingFunction, rowNumberVariable, maxRankPerPartition, partial, hashVariable);
167184
}
168185
}

‎presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,8 @@ public Void visitTopNRowNumber(TopNRowNumberNode node, Void context)
387387
{
388388
printNode(node,
389389
"TopNRowNumber",
390-
format("partition by = %s|order by = %s|n = %s",
390+
format("function = %s; partition by = %s|order by = %s|n = %s",
391+
node.getRankingFunction(),
391392
Joiner.on(", ").join(node.getPartitionBy()),
392393
Joiner.on(", ").join(node.getOrderingScheme().getOrderByVariables()), node.getMaxRowCountPerPartition()),
393394
NODE_COLORS.get(NodeType.WINDOW));

0 commit comments

Comments
 (0)