Skip to content

Commit 45c3305

Browse files
zacw7zacw7
committed
Support native non-equal lookup join planning
This change adds an extractor to traverse the Join plan and get lookup variables in different PlanNode, then stores the lookup variables in LookupJoinNode, which enables index lookup join with non-equal join condition for native execution. Additional changes are made to ensure lookup variables are not pruned by other optimizers.
1 parent 13cd265 commit 45c3305

File tree

9 files changed

+488
-118
lines changed

9 files changed

+488
-118
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -506,7 +506,8 @@ public PlanWithProperties visitIndexJoin(IndexJoinNode node, HashComputationSet
506506
node.getCriteria(),
507507
node.getFilter(),
508508
Optional.of(probeHashVariable),
509-
Optional.of(indexHashVariable)),
509+
Optional.of(indexHashVariable),
510+
node.getLookupVariables()),
510511
allHashVariables);
511512
}
512513

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

Lines changed: 308 additions & 105 deletions
Large diffs are not rendered by default.

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -349,6 +349,8 @@ public PlanNode visitIndexJoin(IndexJoinNode node, RewriteContext<Set<VariableRe
349349
indexInputBuilder.add(node.getIndexHashVariable().get());
350350
}
351351
indexInputBuilder.addAll(expectedFilterInputs);
352+
// Lookup variables must not be pruned.
353+
indexInputBuilder.addAll(node.getLookupVariables());
352354
Set<VariableReferenceExpression> indexInputs = indexInputBuilder.build();
353355

354356
PlanNode probeSource = context.rewrite(node.getProbeSource(), probeInputs);
@@ -364,7 +366,8 @@ public PlanNode visitIndexJoin(IndexJoinNode node, RewriteContext<Set<VariableRe
364366
node.getCriteria(),
365367
node.getFilter(),
366368
node.getProbeHashVariable(),
367-
node.getIndexHashVariable());
369+
node.getIndexHashVariable(),
370+
node.getLookupVariables());
368371
}
369372

370373
@Override

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -660,7 +660,8 @@ public PlanNode visitIndexJoin(IndexJoinNode node, RewriteContext<Void> context)
660660
canonicalizeIndexJoinCriteria(node.getCriteria()),
661661
node.getFilter().map(this::canonicalize),
662662
canonicalize(node.getProbeHashVariable()),
663-
canonicalize(node.getIndexHashVariable()));
663+
canonicalize(node.getIndexHashVariable()),
664+
node.getLookupVariables());
664665
}
665666

666667
@Override

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

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ public class IndexJoinNode
4242
private final Optional<RowExpression> filter;
4343
private final Optional<VariableReferenceExpression> probeHashVariable;
4444
private final Optional<VariableReferenceExpression> indexHashVariable;
45+
private final List<VariableReferenceExpression> lookupVariables;
4546

4647
@JsonCreator
4748
public IndexJoinNode(
@@ -53,7 +54,8 @@ public IndexJoinNode(
5354
@JsonProperty("criteria") List<EquiJoinClause> criteria,
5455
@JsonProperty("filter") Optional<RowExpression> filter,
5556
@JsonProperty("probeHashVariable") Optional<VariableReferenceExpression> probeHashVariable,
56-
@JsonProperty("indexHashVariable") Optional<VariableReferenceExpression> indexHashVariable)
57+
@JsonProperty("indexHashVariable") Optional<VariableReferenceExpression> indexHashVariable,
58+
@JsonProperty("lookupVariables") List<VariableReferenceExpression> lookupVariables)
5759
{
5860
this(sourceLocation,
5961
id,
@@ -64,7 +66,8 @@ public IndexJoinNode(
6466
criteria,
6567
filter,
6668
probeHashVariable,
67-
indexHashVariable);
69+
indexHashVariable,
70+
lookupVariables);
6871
}
6972

7073
public IndexJoinNode(
@@ -77,7 +80,8 @@ public IndexJoinNode(
7780
List<EquiJoinClause> criteria,
7881
Optional<RowExpression> filter,
7982
Optional<VariableReferenceExpression> probeHashVariable,
80-
Optional<VariableReferenceExpression> indexHashVariable)
83+
Optional<VariableReferenceExpression> indexHashVariable,
84+
List<VariableReferenceExpression> lookupVariables)
8185
{
8286
super(sourceLocation, id, statsEquivalentPlanNode);
8387
this.type = requireNonNull(type, "type is null");
@@ -87,6 +91,7 @@ public IndexJoinNode(
8791
this.filter = requireNonNull(filter, "filter is null");
8892
this.probeHashVariable = requireNonNull(probeHashVariable, "probeHashVariable is null");
8993
this.indexHashVariable = requireNonNull(indexHashVariable, "indexHashVariable is null");
94+
this.lookupVariables = requireNonNull(lookupVariables, "lookupVariables is null");
9095
}
9196

9297
@JsonProperty
@@ -131,6 +136,12 @@ public Optional<VariableReferenceExpression> getIndexHashVariable()
131136
return indexHashVariable;
132137
}
133138

139+
@JsonProperty
140+
public List<VariableReferenceExpression> getLookupVariables()
141+
{
142+
return lookupVariables;
143+
}
144+
134145
@Override
135146
public List<PlanNode> getSources()
136147
{
@@ -166,7 +177,8 @@ public PlanNode replaceChildren(List<PlanNode> newChildren)
166177
criteria,
167178
filter,
168179
probeHashVariable,
169-
indexHashVariable);
180+
indexHashVariable,
181+
lookupVariables);
170182
}
171183

172184
@Override
@@ -182,7 +194,8 @@ public PlanNode assignStatsEquivalentPlanNode(Optional<PlanNode> statsEquivalent
182194
criteria,
183195
filter,
184196
probeHashVariable,
185-
indexHashVariable);
197+
indexHashVariable,
198+
lookupVariables);
186199
}
187200

188201
public static class EquiJoinClause

‎presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -472,9 +472,7 @@ public Void visitIndexJoin(IndexJoinNode node, Set<VariableReferenceExpression>
472472
checkArgument(indexSourceInputs.contains(clause.getIndex()), "Index variable from index join clause (%s) not in index source (%s)", clause.getIndex(), node.getIndexSource().getOutputVariables());
473473
}
474474

475-
Set<VariableReferenceExpression> lookupVariables = node.getCriteria().stream()
476-
.map(IndexJoinNode.EquiJoinClause::getIndex)
477-
.collect(toImmutableSet());
475+
Set<VariableReferenceExpression> lookupVariables = ImmutableSet.copyOf(node.getLookupVariables());
478476
Map<VariableReferenceExpression, VariableReferenceExpression> trace = IndexKeyTracer.trace(node.getIndexSource(), lookupVariables);
479477
checkArgument(!trace.isEmpty() && lookupVariables.containsAll(trace.keySet()),
480478
"Index lookup symbols are not traceable to index source: %s",

‎presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -865,7 +865,8 @@ public PlanNode indexJoin(JoinType type,
865865
criteria,
866866
filter,
867867
Optional.empty(),
868-
Optional.empty());
868+
Optional.empty(),
869+
index.getOutputVariables());
869870
}
870871

871872
public CteProducerNode cteProducerNode(String ctename,

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,8 @@ public void testIndexJoinNodePruning()
191191
output(
192192
indexJoin(
193193
strictTableScan("lineitem", ImmutableMap.of("partkey", "partkey", "suppkey", "suppkey")),
194-
strictIndexSource("orders", ImmutableMap.of("custkey", "custkey", "orderkey", "orderkey")))));
194+
strictIndexSource("orders",
195+
ImmutableMap.of("custkey", "custkey", "orderkey", "orderkey", "orderstatus", "orderstatus", "totalprice", "totalprice")))));
195196
}
196197

197198
private OptimizerAssert assertRuleApplication()

‎presto-tests/src/test/java/com/facebook/presto/tests/TestNativeIndexJoinLogicalPlanner.java‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.tests;
15+
16+
import com.facebook.presto.Session;
17+
import com.facebook.presto.testing.QueryRunner;
18+
import com.facebook.presto.tests.tpch.IndexedTpchPlugin;
19+
import com.google.common.collect.ImmutableList;
20+
import org.testng.annotations.Test;
21+
22+
import java.util.List;
23+
24+
import static com.facebook.presto.SystemSessionProperties.NATIVE_EXECUTION_ENABLED;
25+
import static com.facebook.presto.SystemSessionProperties.OPTIMIZE_HASH_GENERATION;
26+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree;
27+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter;
28+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.indexJoin;
29+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.indexSource;
30+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project;
31+
import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableScan;
32+
import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
33+
import static com.facebook.presto.tests.AbstractTestIndexedQueries.INDEX_SPEC;
34+
import static com.facebook.presto.tpch.TpchMetadata.TINY_SCHEMA_NAME;
35+
36+
public class TestNativeIndexJoinLogicalPlanner
37+
extends AbstractTestQueryFramework
38+
{
39+
public static final List<String> SUPPORTED_JOIN_TYPES = ImmutableList.of("INNER", "LEFT");
40+
41+
@Override
42+
protected QueryRunner createQueryRunner()
43+
throws Exception
44+
{
45+
Session session = testSessionBuilder()
46+
.setCatalog("tpch_indexed")
47+
.setSchema(TINY_SCHEMA_NAME)
48+
.setSystemProperty(NATIVE_EXECUTION_ENABLED, "true")
49+
.setSystemProperty(OPTIMIZE_HASH_GENERATION, "false")
50+
.build();
51+
52+
DistributedQueryRunner queryRunner = new DistributedQueryRunner.Builder(session)
53+
.setNodeCount(1)
54+
.build();
55+
56+
queryRunner.installPlugin(new IndexedTpchPlugin(INDEX_SPEC));
57+
queryRunner.createCatalog("tpch_indexed", "tpch_indexed");
58+
return queryRunner;
59+
}
60+
61+
@Test
62+
public void testBasicIndexJoin()
63+
{
64+
for (String joinType : SUPPORTED_JOIN_TYPES) {
65+
assertPlan("" +
66+
"SELECT *\n" +
67+
"FROM (\n" +
68+
" SELECT *\n" +
69+
" FROM lineitem\n" +
70+
" WHERE partkey % 8 = 0) l\n" +
71+
joinType + " JOIN orders o\n" +
72+
" ON l.orderkey = o.orderkey",
73+
anyTree(indexJoin(
74+
filter(tableScan("lineitem")),
75+
indexSource("orders"))));
76+
77+
assertPlan("" +
78+
"SELECT *\n" +
79+
"FROM (\n" +
80+
" SELECT CASE WHEN suppkey % 2 = 0 THEN 'F' ELSE 'O' END AS orderstatus, *\n" +
81+
" FROM lineitem\n" +
82+
" WHERE partkey % 8 = 0) l\n" +
83+
joinType + " JOIN orders o\n" +
84+
" ON l.orderkey = o.orderkey\n" +
85+
" AND l.orderstatus = o.orderstatus",
86+
anyTree(indexJoin(
87+
project(filter(tableScan("lineitem"))),
88+
indexSource("orders"))));
89+
}
90+
}
91+
92+
@Test
93+
public void testNonEqualIndexJoin()
94+
{
95+
for (String joinType : SUPPORTED_JOIN_TYPES) {
96+
assertPlan("" +
97+
"SELECT *\n" +
98+
"FROM (\n" +
99+
" SELECT *\n" +
100+
" FROM lineitem\n" +
101+
" WHERE partkey % 8 = 0) l\n" +
102+
joinType + " JOIN orders o\n" +
103+
" ON l.orderkey = o.orderkey" +
104+
" AND o.custkey BETWEEN 1 AND l.partkey",
105+
anyTree(indexJoin(
106+
filter(tableScan("lineitem")),
107+
indexSource("orders"))));
108+
109+
assertPlan("" +
110+
"SELECT *\n" +
111+
"FROM (\n" +
112+
" SELECT *\n" +
113+
" FROM lineitem\n" +
114+
" WHERE partkey % 8 = 0) l\n" +
115+
joinType + " JOIN orders o\n" +
116+
" ON l.orderkey = o.orderkey" +
117+
" AND CONTAINS(ARRAY[1, l.partkey, 3], o.custkey)",
118+
anyTree(indexJoin(
119+
filter(tableScan("lineitem")),
120+
indexSource("orders"))));
121+
122+
assertPlan("" +
123+
"SELECT *\n" +
124+
"FROM (\n" +
125+
" SELECT *\n" +
126+
" FROM lineitem\n" +
127+
" WHERE partkey % 8 = 0) l\n" +
128+
joinType + " JOIN orders o\n" +
129+
" ON l.orderkey = o.orderkey" +
130+
" AND o.custkey BETWEEN 1 AND 100",
131+
anyTree(indexJoin(
132+
filter(tableScan("lineitem")),
133+
filter(indexSource("orders")))));
134+
135+
assertPlan("" +
136+
"SELECT *\n" +
137+
"FROM (\n" +
138+
" SELECT *\n" +
139+
" FROM lineitem\n" +
140+
" WHERE partkey % 8 = 0) l\n" +
141+
joinType + " JOIN orders o\n" +
142+
" ON l.orderkey = o.orderkey" +
143+
" AND CONTAINS(ARRAY[1, 2, 3], o.custkey)",
144+
anyTree(indexJoin(
145+
filter(tableScan("lineitem")),
146+
filter(indexSource("orders")))));
147+
}
148+
}
149+
}

0 commit comments

Comments
 (0)