Skip to content

Commit c7a1a34

Browse files
bestbeforetodaybestbeforetoday
authored
feat(isthmus): support subquery nested in project (#530)
Previously a subquery nested within a project was explicitly disallowed when converting from Calcite to Substrait. This is because outer references were not correctly resolved within nested subqueries. This change tracks outer references so they can be successfully resolved in subqueries nested within a project. Signed-off-by: Mark S. Lewis <Mark.S.Lewis@outlook.com>
1 parent c2eb5f7 commit c7a1a34

3 files changed

Lines changed: 45 additions & 65 deletions

File tree

‎isthmus/src/main/java/io/substrait/isthmus/OuterReferenceResolver.java‎

Lines changed: 18 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
package io.substrait.isthmus;
22

3-
import static org.apache.calcite.rex.RexUtil.SubQueryFinder.containsSubQuery;
4-
53
import java.util.HashMap;
64
import java.util.IdentityHashMap;
75
import java.util.Map;
@@ -15,15 +13,16 @@
1513
import org.apache.calcite.rex.RexNode;
1614
import org.apache.calcite.rex.RexShuttle;
1715
import org.apache.calcite.rex.RexSubQuery;
16+
import org.apache.calcite.rex.RexUtil.SubQueryCollector;
1817

1918
/** Resolve correlated variable and get Depth map for RexFieldAccess */
2019
// See OuterReferenceResolver.md for explanation how the Depth map is computed.
2120
public class OuterReferenceResolver extends RelNodeVisitor<RelNode, RuntimeException> {
2221

23-
private Map<CorrelationId, Integer> nestedDepth;
24-
private Map<RexFieldAccess, Integer> fieldAccessDepthMap;
22+
private final Map<CorrelationId, Integer> nestedDepth;
23+
private final Map<RexFieldAccess, Integer> fieldAccessDepthMap;
2524

26-
private RexVisitor rexVisitor = new RexVisitor(this);
25+
private final RexVisitor rexVisitor = new RexVisitor(this);
2726

2827
public OuterReferenceResolver() {
2928
nestedDepth = new HashMap<>();
@@ -45,9 +44,7 @@ public Map<RexFieldAccess, Integer> getFieldAccessDepthMap() {
4544
@Override
4645
public RelNode visit(Filter filter) throws RuntimeException {
4746
for (CorrelationId id : filter.getVariablesSet()) {
48-
if (!nestedDepth.containsKey(id)) {
49-
nestedDepth.put(id, 0);
50-
}
47+
nestedDepth.putIfAbsent(id, 0);
5148
}
5249
filter.getCondition().accept(rexVisitor);
5350
return super.visit(filter);
@@ -56,25 +53,19 @@ public RelNode visit(Filter filter) throws RuntimeException {
5653
@Override
5754
public RelNode visit(Correlate correlate) throws RuntimeException {
5855
for (CorrelationId id : correlate.getVariablesSet()) {
59-
if (!nestedDepth.containsKey(id)) {
60-
nestedDepth.put(id, 0);
61-
}
56+
nestedDepth.putIfAbsent(id, 0);
6257
}
6358

6459
apply(correlate.getLeft());
6560

6661
// Correlated join is a special case. The right-rel is a correlated sub-query but not a REX. So,
6762
// the RexVisitor cannot be applied to it to correctly compute the depth map. Hence, we need to
6863
// manually compute the depth map for the right-rel.
69-
for (Map.Entry<CorrelationId, Integer> entry : nestedDepth.entrySet()) {
70-
nestedDepth.put(entry.getKey(), entry.getValue() + 1);
71-
}
64+
nestedDepth.replaceAll((k, v) -> v + 1);
7265

7366
apply(correlate.getRight()); // look inside sub-queries
7467

75-
for (Map.Entry<CorrelationId, Integer> entry : nestedDepth.entrySet()) {
76-
nestedDepth.put(entry.getKey(), entry.getValue() - 1);
77-
}
68+
nestedDepth.replaceAll((k, v) -> v - 1);
7869

7970
return correlate;
8071
}
@@ -89,14 +80,18 @@ public RelNode visitOther(RelNode other) throws RuntimeException {
8980

9081
@Override
9182
public RelNode visit(Project project) throws RuntimeException {
92-
if (containsSubQuery(project)) {
93-
throw new UnsupportedOperationException(
94-
"Unsupported subquery nested in Project relational operator : " + project);
83+
for (CorrelationId id : project.getVariablesSet()) {
84+
nestedDepth.putIfAbsent(id, 0);
9585
}
86+
87+
for (RexSubQuery subQuery : SubQueryCollector.collect(project)) {
88+
subQuery.accept(rexVisitor);
89+
}
90+
9691
return super.visit(project);
9792
}
9893

99-
private class RexVisitor extends RexShuttle {
94+
private static class RexVisitor extends RexShuttle {
10095
final OuterReferenceResolver referenceResolver;
10196

10297
RexVisitor(OuterReferenceResolver referenceResolver) {
@@ -105,15 +100,11 @@ private class RexVisitor extends RexShuttle {
105100

106101
@Override
107102
public RexNode visitSubQuery(RexSubQuery subQuery) {
108-
for (Map.Entry<CorrelationId, Integer> entry : referenceResolver.nestedDepth.entrySet()) {
109-
referenceResolver.nestedDepth.put(entry.getKey(), entry.getValue() + 1);
110-
}
103+
referenceResolver.nestedDepth.replaceAll((k, v) -> v + 1);
111104

112105
referenceResolver.apply(subQuery.rel); // look inside sub-queries
113106

114-
for (Map.Entry<CorrelationId, Integer> entry : referenceResolver.nestedDepth.entrySet()) {
115-
referenceResolver.nestedDepth.put(entry.getKey(), entry.getValue() - 1);
116-
}
107+
referenceResolver.nestedDepth.replaceAll((k, v) -> v - 1);
117108
return subQuery;
118109
}
119110

‎isthmus/src/test/java/io/substrait/isthmus/SubqueryPlanTest.java‎

Lines changed: 26 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,22 @@
11
package io.substrait.isthmus;
22

3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertSame;
35
import static org.junit.jupiter.api.Assertions.assertTrue;
46

57
import com.google.protobuf.util.JsonFormat;
68
import io.substrait.proto.Expression;
9+
import io.substrait.proto.Expression.Subquery.SetPredicate.PredicateOp;
710
import io.substrait.proto.FilterRel;
811
import io.substrait.proto.Plan;
912
import java.io.IOException;
1013
import org.apache.calcite.sql.parser.SqlParseException;
11-
import org.junit.jupiter.api.Assertions;
1214
import org.junit.jupiter.api.Test;
1315

1416
public class SubqueryPlanTest extends PlanTestBase {
1517
// TODO: Add a roundtrip test once the ProtoRelConverter is committed and updated to support
1618
// subqueries
19+
1720
@Test
1821
public void existsCorrelatedSubquery() throws SqlParseException {
1922
SqlToSubstrait s = new SqlToSubstrait();
@@ -34,9 +37,7 @@ public void existsCorrelatedSubquery() throws SqlParseException {
3437
.getSubquery();
3538

3639
assertTrue(subquery.hasSetPredicate());
37-
assertTrue(
38-
subquery.getSetPredicate().getPredicateOp()
39-
== Expression.Subquery.SetPredicate.PredicateOp.PREDICATE_OP_EXISTS);
40+
assertSame(PredicateOp.PREDICATE_OP_EXISTS, subquery.getSetPredicate().getPredicateOp());
4041

4142
FilterRel setPredicateFilter =
4243
subquery
@@ -52,8 +53,8 @@ public void existsCorrelatedSubquery() throws SqlParseException {
5253
.getValue()
5354
.getSelection(); // l_orderkey
5455

55-
assertTrue(correlatedCol.getDirectReference().getStructField().getField() == 0);
56-
assertTrue(correlatedCol.getOuterReference().getStepsOut() == 1);
56+
assertEquals(0, correlatedCol.getDirectReference().getStructField().getField());
57+
assertEquals(1, correlatedCol.getOuterReference().getStepsOut());
5758
}
5859

5960
@Test
@@ -85,9 +86,7 @@ public void uniqueCorrelatedSubquery() throws IOException, SqlParseException {
8586
.getFilter(); // unique (select ... from orders where o_orderkey = l_orderkey)
8687

8788
assertTrue(subquery.hasSetPredicate());
88-
assertTrue(
89-
subquery.getSetPredicate().getPredicateOp()
90-
== Expression.Subquery.SetPredicate.PredicateOp.PREDICATE_OP_UNIQUE);
89+
assertSame(PredicateOp.PREDICATE_OP_UNIQUE, subquery.getSetPredicate().getPredicateOp());
9190

9291
Expression.FieldReference correlatedCol =
9392
setPredicateFilter
@@ -97,8 +96,8 @@ public void uniqueCorrelatedSubquery() throws IOException, SqlParseException {
9796
.getValue()
9897
.getSelection(); // l_orderkey
9998

100-
assertTrue(correlatedCol.getDirectReference().getStructField().getField() == 0);
101-
assertTrue(correlatedCol.getOuterReference().getStepsOut() == 1);
99+
assertEquals(0, correlatedCol.getDirectReference().getStructField().getField());
100+
assertEquals(1, correlatedCol.getOuterReference().getStepsOut());
102101
}
103102

104103
@Test
@@ -135,8 +134,8 @@ public void inPredicateCorrelatedSubQuery() throws IOException, SqlParseExceptio
135134
.getValue()
136135
.getSelection(); // l_partkey
137136

138-
assertTrue(correlatedCol.getDirectReference().getStructField().getField() == 1);
139-
assertTrue(correlatedCol.getOuterReference().getStepsOut() == 1);
137+
assertEquals(1, correlatedCol.getDirectReference().getStructField().getField());
138+
assertEquals(1, correlatedCol.getOuterReference().getStepsOut());
140139
}
141140

142141
@Test
@@ -175,8 +174,8 @@ public void notInPredicateCorrelatedSubquery() throws IOException, SqlParseExcep
175174
.getValue()
176175
.getSelection(); // l_partkey
177176

178-
assertTrue(correlatedCol.getDirectReference().getStructField().getField() == 1);
179-
assertTrue(correlatedCol.getOuterReference().getStepsOut() == 1);
177+
assertEquals(1, correlatedCol.getDirectReference().getStructField().getField());
178+
assertEquals(1, correlatedCol.getOuterReference().getStepsOut());
180179
}
181180

182181
@Test
@@ -207,9 +206,7 @@ public void existsNestedCorrelatedSubquery() throws IOException, SqlParseExcepti
207206
.getSubquery();
208207

209208
assertTrue(outer_subquery.hasSetPredicate());
210-
assertTrue(
211-
outer_subquery.getSetPredicate().getPredicateOp()
212-
== Expression.Subquery.SetPredicate.PredicateOp.PREDICATE_OP_EXISTS);
209+
assertSame(PredicateOp.PREDICATE_OP_EXISTS, outer_subquery.getSetPredicate().getPredicateOp());
213210

214211
FilterRel exists_filter =
215212
outer_subquery
@@ -221,9 +218,7 @@ public void existsNestedCorrelatedSubquery() throws IOException, SqlParseExcepti
221218
exists_filter.getCondition().getScalarFunction().getArguments(1).getValue().getSubquery();
222219
assertTrue(inner_subquery.hasSetPredicate());
223220

224-
assertTrue(
225-
inner_subquery.getSetPredicate().getPredicateOp()
226-
== Expression.Subquery.SetPredicate.PredicateOp.PREDICATE_OP_UNIQUE);
221+
assertSame(PredicateOp.PREDICATE_OP_UNIQUE, inner_subquery.getSetPredicate().getPredicateOp());
227222

228223
Expression inner_subquery_condition =
229224
inner_subquery
@@ -251,17 +246,17 @@ public void existsNestedCorrelatedSubquery() throws IOException, SqlParseExcepti
251246
.getArguments(1)
252247
.getValue()
253248
.getSelection(); // p.p_partkey
254-
assertTrue(correlatedCol1.getDirectReference().getStructField().getField() == 0);
255-
assertTrue(correlatedCol1.getOuterReference().getStepsOut() == 2);
249+
assertEquals(0, correlatedCol1.getDirectReference().getStructField().getField());
250+
assertEquals(2, correlatedCol1.getOuterReference().getStepsOut());
256251

257252
Expression.FieldReference correlatedCol2 =
258253
inner_subquery_cond2
259254
.getScalarFunction()
260255
.getArguments(1)
261256
.getValue()
262257
.getSelection(); // l.l_suppkey
263-
assertTrue(correlatedCol2.getDirectReference().getStructField().getField() == 2);
264-
assertTrue(correlatedCol2.getOuterReference().getStepsOut() == 1);
258+
assertEquals(2, correlatedCol2.getDirectReference().getStructField().getField());
259+
assertEquals(1, correlatedCol2.getOuterReference().getStepsOut());
265260
}
266261

267262
@Test
@@ -326,27 +321,22 @@ public void nestedScalarCorrelatedSubquery() throws IOException, SqlParseExcepti
326321
.getArguments(1)
327322
.getValue()
328323
.getSelection(); // p.p_partkey
329-
assertTrue(correlatedCol1.getDirectReference().getStructField().getField() == 0);
330-
assertTrue(correlatedCol1.getOuterReference().getStepsOut() == 2);
324+
assertEquals(0, correlatedCol1.getDirectReference().getStructField().getField());
325+
assertEquals(2, correlatedCol1.getOuterReference().getStepsOut());
331326

332327
Expression.FieldReference correlatedCol2 =
333328
inner_subquery_cond2
334329
.getScalarFunction()
335330
.getArguments(1)
336331
.getValue()
337332
.getSelection(); // l.l_suppkey
338-
assertTrue(correlatedCol2.getDirectReference().getStructField().getField() == 2);
339-
assertTrue(correlatedCol2.getOuterReference().getStepsOut() == 1);
333+
assertEquals(2, correlatedCol2.getDirectReference().getStructField().getField());
334+
assertEquals(1, correlatedCol2.getOuterReference().getStepsOut());
340335
}
341336

342337
@Test
343-
public void correlatedScalarSubQInSelect() throws IOException {
344-
SqlToSubstrait s = new SqlToSubstrait();
338+
public void correlatedScalarSubQueryInSelect() throws Exception {
345339
String sql = asString("subquery/nested_scalar_subquery_in_select.sql");
346-
Assertions.assertThrows(
347-
UnsupportedOperationException.class,
348-
() -> {
349-
s.convert(sql, TPCH_CATALOG);
350-
});
340+
assertSqlSubstraitRelRoundTrip(sql);
351341
}
352342
}

‎isthmus/src/test/java/io/substrait/isthmus/TpcdsQueryTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,11 @@
1414
/** TPC-DS test to convert SQL to Substrait and then convert those plans back to SQL. */
1515
public class TpcdsQueryTest extends PlanTestBase {
1616
private static final Set<Integer> alternateForms = Set.of(27, 36, 70, 86);
17-
private static final Set<Integer> toSubstraitExclusions = Set.of(9);
1817
private static final Set<Integer> fromSubstraitPojoExclusions = Set.of(1, 30, 81);
1918
private static final Set<Integer> fromSubstraitProtoExclusions = Set.of(1, 30, 81);
2019

2120
static IntStream testCases() {
22-
return IntStream.rangeClosed(1, 99).filter(n -> !toSubstraitExclusions.contains(n));
21+
return IntStream.rangeClosed(1, 99);
2322
}
2423

2524
/**

0 commit comments

Comments
 (0)