Skip to content

Commit 418742f

Browse files
nielspardonnielspardon
authored
feat(isthmus): introduce SqlToSubstrait#convert method (#473)
SqlToSubstrait#convert converts a SQL query to a POJO plan As part of this change SqlToSubstrait#execute has been deprecated, as that method returned a PROTO plan
1 parent 0532c18 commit 418742f

12 files changed

Lines changed: 126 additions & 99 deletions

File tree

‎isthmus-cli/src/main/java/io/substrait/isthmus/cli/IsthmusEntryPoint.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import io.substrait.isthmus.SqlExpressionToSubstrait;
1111
import io.substrait.isthmus.SqlToSubstrait;
1212
import io.substrait.isthmus.sql.SubstraitCreateStatementParser;
13+
import io.substrait.plan.PlanProtoConverter;
1314
import io.substrait.proto.ExtendedExpression;
1415
import io.substrait.proto.Plan;
1516
import java.io.IOException;
@@ -94,7 +95,7 @@ public Integer call() throws Exception {
9495
Prepare.CatalogReader catalog =
9596
SubstraitCreateStatementParser.processCreateStatementsToCatalog(
9697
createStatements.toArray(String[]::new));
97-
Plan plan = converter.execute(sql, catalog);
98+
Plan plan = new PlanProtoConverter().toProto(converter.convert(sql, catalog));
9899
printMessage(plan);
99100
}
100101
return 0;

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

Lines changed: 35 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
import com.google.common.annotations.VisibleForTesting;
44
import io.substrait.isthmus.sql.SubstraitSqlValidator;
55
import io.substrait.plan.ImmutablePlan.Builder;
6+
import io.substrait.plan.Plan;
67
import io.substrait.plan.Plan.Version;
78
import io.substrait.plan.PlanProtoConverter;
8-
import io.substrait.proto.Plan;
99
import java.util.List;
1010
import org.apache.calcite.plan.hep.HepPlanner;
1111
import org.apache.calcite.plan.hep.HepProgram;
@@ -31,35 +31,51 @@ public SqlToSubstrait(FeatureBoard features) {
3131
super(features);
3232
}
3333

34-
public Plan execute(String sql, Prepare.CatalogReader catalogReader) throws SqlParseException {
35-
SqlValidator validator = new SubstraitSqlValidator(catalogReader);
36-
return executeInner(sql, validator, catalogReader);
37-
}
38-
39-
List<RelRoot> sqlToRelNode(String sql, Prepare.CatalogReader catalogReader)
34+
/**
35+
* Converts a SQL statements string into a Substrait proto {@link io.substrait.proto.Plan}.
36+
*
37+
* @param sql the SQL statements string containing one more SQL statements
38+
* @param catalogReader the {@link Prepare.CatalogReader} for finding tables/views referenced in
39+
* the SQL statements string
40+
* @return the Substrait proto {@link io.substrait.proto.Plan}
41+
* @throws SqlParseException if there is an error while parsing the SQL statements string
42+
* @deprecated use {@link #convert(String, org.apache.calcite.prepare.Prepare.CatalogReader)}
43+
* instead to get a {@link Plan} and convert that to a {@link io.substrait.proto.Plan} using
44+
* {@link PlanProtoConverter#toProto(Plan)}
45+
*/
46+
@Deprecated
47+
public io.substrait.proto.Plan execute(String sql, Prepare.CatalogReader catalogReader)
4048
throws SqlParseException {
41-
SqlValidator validator = new SubstraitSqlValidator(catalogReader);
42-
return sqlToRelNode(sql, validator, catalogReader);
49+
PlanProtoConverter planToProto = new PlanProtoConverter();
50+
51+
return planToProto.toProto(convert(sql, catalogReader));
4352
}
4453

45-
private Plan executeInner(String sql, SqlValidator validator, Prepare.CatalogReader catalogReader)
46-
throws SqlParseException {
54+
/**
55+
* Converts a SQL statements string into a Substrait {@link Plan}.
56+
*
57+
* @param sql the SQL statements string containing one more SQL statements
58+
* @param catalogReader the {@link Prepare.CatalogReader} for finding tables/views referenced in
59+
* the SQL statements string
60+
* @return the Substrait {@link Plan}
61+
* @throws SqlParseException if there is an error while parsing the SQL statements string
62+
*/
63+
public Plan convert(String sql, Prepare.CatalogReader catalogReader) throws SqlParseException {
4764
Builder builder = io.substrait.plan.Plan.builder();
4865
builder.version(Version.builder().from(Version.DEFAULT_VERSION).producer("isthmus").build());
4966

5067
// TODO: consider case in which one sql passes conversion while others don't
51-
sqlToRelNode(sql, validator, catalogReader).stream()
68+
sqlToRelNode(sql, catalogReader).stream()
5269
.map(root -> SubstraitRelVisitor.convert(root, EXTENSION_COLLECTION, featureBoard))
5370
.forEach(root -> builder.addRoots(root));
5471

55-
PlanProtoConverter planToProto = new PlanProtoConverter();
56-
57-
return planToProto.toProto(builder.build());
72+
return builder.build();
5873
}
5974

60-
private List<RelRoot> sqlToRelNode(
61-
String sql, SqlValidator validator, Prepare.CatalogReader catalogReader)
75+
@VisibleForTesting
76+
List<RelRoot> sqlToRelNode(String sql, Prepare.CatalogReader catalogReader)
6277
throws SqlParseException {
78+
SqlValidator validator = new SubstraitSqlValidator(catalogReader);
6379
SqlParser parser = SqlParser.create(sql, parserConfig);
6480
SqlNodeList parsedList = parser.parseStmtList();
6581
SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);
@@ -70,8 +86,7 @@ private List<RelRoot> sqlToRelNode(
7086
return roots;
7187
}
7288

73-
@VisibleForTesting
74-
SqlToRelConverter createSqlToRelConverter(
89+
protected SqlToRelConverter createSqlToRelConverter(
7590
SqlValidator validator, Prepare.CatalogReader catalogReader) {
7691
SqlToRelConverter converter =
7792
new SqlToRelConverter(
@@ -84,8 +99,7 @@ SqlToRelConverter createSqlToRelConverter(
8499
return converter;
85100
}
86101

87-
@VisibleForTesting
88-
static RelRoot getBestExpRelRoot(SqlToRelConverter converter, SqlNode parsed) {
102+
protected RelRoot getBestExpRelRoot(SqlToRelConverter converter, SqlNode parsed) {
89103
RelRoot root = converter.convertQuery(parsed, true, true);
90104
{
91105
// RelBuilder seems to implicitly use the rule below,

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public void lateralJoinQuery() throws SqlParseException {
6262
SqlToSubstrait sE2E = new SqlToSubstrait();
6363
Assertions.assertThrows(
6464
UnsupportedOperationException.class,
65-
() -> sE2E.execute(sql, TPCDS_CATALOG),
65+
() -> sE2E.convert(sql, TPCDS_CATALOG),
6666
"Lateral join is not supported");
6767
}
6868

@@ -83,7 +83,7 @@ public void outerApplyQuery() throws SqlParseException {
8383
// TODO validate end to end conversion
8484
Assertions.assertThrows(
8585
UnsupportedOperationException.class,
86-
() -> new SqlToSubstrait().execute(sql, TPCDS_CATALOG),
86+
() -> new SqlToSubstrait().convert(sql, TPCDS_CATALOG),
8787
"APPLY is not supported");
8888
}
8989

@@ -123,7 +123,7 @@ public void nestedApplyJoinQuery() throws SqlParseException {
123123
// TODO validate end to end conversion
124124
Assertions.assertThrows(
125125
UnsupportedOperationException.class,
126-
() -> new SqlToSubstrait().execute(sql, TPCDS_CATALOG),
126+
() -> new SqlToSubstrait().convert(sql, TPCDS_CATALOG),
127127
"APPLY is not supported");
128128
}
129129

@@ -138,7 +138,7 @@ public void crossApplyQuery() throws SqlParseException {
138138
// TODO validate end to end conversion
139139
Assertions.assertThrows(
140140
UnsupportedOperationException.class,
141-
() -> new SqlToSubstrait().execute(sql, TPCDS_CATALOG),
141+
() -> new SqlToSubstrait().convert(sql, TPCDS_CATALOG),
142142
"APPLY is not supported");
143143
}
144144
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,6 @@ public void tpcds(int query) throws Exception {
1414
String values = asString("keyconstraints_schema.sql");
1515
Prepare.CatalogReader catalog =
1616
SubstraitCreateStatementParser.processCreateStatementsToCatalog(values);
17-
s.execute(asString(String.format("tpcds/queries/%02d.sql", query)), catalog);
17+
s.convert(asString(String.format("tpcds/queries/%02d.sql", query)), catalog);
1818
}
1919
}

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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,10 @@ void preserveNamesFromSql() throws Exception {
2525
String query = "SELECT \"a\", \"B\" FROM foo GROUP BY a, b";
2626
List<String> expectedNames = List.of("a", "B");
2727

28-
List<org.apache.calcite.rel.RelRoot> calciteRelRoots = s.sqlToRelNode(query, catalogReader);
29-
assertEquals(1, calciteRelRoots.size());
28+
Plan plan = s.convert(query, catalogReader);
29+
assertEquals(1, plan.getRoots().size());
3030

31-
org.apache.calcite.rel.RelRoot calciteRelRoot1 = calciteRelRoots.get(0);
32-
assertEquals(expectedNames, calciteRelRoot1.validatedRowType.getFieldNames());
33-
34-
io.substrait.plan.Plan.Root substraitRelRoot =
35-
SubstraitRelVisitor.convert(calciteRelRoot1, EXTENSION_COLLECTION);
31+
io.substrait.plan.Plan.Root substraitRelRoot = plan.getRoots().get(0);
3632
assertEquals(expectedNames, substraitRelRoot.getNames());
3733

3834
org.apache.calcite.rel.RelRoot calciteRelRoot2 = substraitToCalcite.convert(substraitRelRoot);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ private void test(Table table, String query, String expectedExpressionText)
6060
final Schema schema = new SubstraitSchema(Map.of("my_table", table));
6161
final CalciteCatalogReader catalog = schemaToCatalog("nested", schema);
6262
final SqlToSubstrait sqlToSubstrait = new SqlToSubstrait();
63-
Plan plan = sqlToSubstrait.execute(query, catalog);
63+
Plan plan = toProto(sqlToSubstrait.convert(query, catalog));
6464
Expression obtainedExpression =
6565
plan.getRelations(0).getRoot().getInput().getProject().getExpressions(0);
6666
Expression expectedExpression = TextFormat.parse(expectedExpressionText, Expression.class);

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

Lines changed: 40 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -87,24 +87,26 @@ protected Plan assertProtoPlanRoundrip(String query, SqlToSubstrait s, String cr
8787
protected Plan assertProtoPlanRoundrip(
8888
String query, SqlToSubstrait s, Prepare.CatalogReader catalogReader)
8989
throws SqlParseException {
90-
io.substrait.proto.Plan protoPlan1 = s.execute(query, catalogReader);
91-
Plan plan = new ProtoPlanConverter(extensions).from(protoPlan1);
92-
io.substrait.proto.Plan protoPlan2 = new PlanProtoConverter().toProto(plan);
90+
Plan plan1 = s.convert(query, catalogReader);
91+
io.substrait.proto.Plan protoPlan1 = toProto(plan1);
92+
93+
Plan plan2 = new ProtoPlanConverter(extensions).from(protoPlan1);
94+
io.substrait.proto.Plan protoPlan2 = toProto(plan2);
9395
assertEquals(protoPlan1, protoPlan2);
94-
List<RelRoot> rootRels = s.sqlToRelNode(query, catalogReader);
95-
assertEquals(rootRels.size(), plan.getRoots().size());
96-
for (int i = 0; i < rootRels.size(); i++) {
97-
Plan.Root rootRel = SubstraitRelVisitor.convert(rootRels.get(i), extensions);
96+
97+
assertEquals(plan1.getRoots().size(), plan2.getRoots().size());
98+
for (int i = 0; i < plan1.getRoots().size(); i++) {
9899
assertEquals(
99-
rootRel.getInput().getRecordType(), plan.getRoots().get(i).getInput().getRecordType());
100+
plan1.getRoots().get(i).getInput().getRecordType(),
101+
plan2.getRoots().get(i).getInput().getRecordType());
100102
}
101-
return plan;
103+
104+
return plan2;
102105
}
103106

104107
protected void assertPlanRoundtrip(Plan plan) {
105-
io.substrait.proto.Plan protoPlan1 = new PlanProtoConverter().toProto(plan);
106-
io.substrait.proto.Plan protoPlan2 =
107-
new PlanProtoConverter().toProto(new ProtoPlanConverter().from(protoPlan1));
108+
io.substrait.proto.Plan protoPlan1 = toProto(plan);
109+
io.substrait.proto.Plan protoPlan2 = toProto(new ProtoPlanConverter().from(protoPlan1));
108110
assertEquals(protoPlan1, protoPlan2);
109111
}
110112

@@ -129,13 +131,11 @@ protected RelRoot assertSqlSubstraitRelRoundTrip(
129131

130132
SqlToSubstrait s = new SqlToSubstrait();
131133

132-
// 1. SQL -> Calcite RelRoot
133-
List<RelRoot> relRoots = s.sqlToRelNode(query, catalogReader);
134-
assertEquals(1, relRoots.size());
135-
RelRoot relRoot1 = relRoots.get(0);
134+
// 1. SQL -> Substrait Plan
135+
Plan plan1 = s.convert(query, catalogReader);
136136

137-
// 2. Calcite RelRoot -> Substrait Rel
138-
Plan.Root pojo1 = SubstraitRelVisitor.convert(relRoot1, extensions);
137+
// 2. Substrait Plan -> Substrait Rel
138+
Plan.Root pojo1 = plan1.getRoots().get(0);
139139

140140
// 3. Substrait Rel -> Calcite RelNode
141141
RelRoot relRoot2 = substraitToCalcite.convert(pojo1);
@@ -178,37 +178,36 @@ protected void assertFullRoundTrip(String sqlQuery, Prepare.CatalogReader catalo
178178
SqlToSubstrait sqlConverter = new SqlToSubstrait();
179179
ExtensionCollector extensionCollector = new ExtensionCollector();
180180

181-
// SQL -> Calcite 1
182-
List<RelRoot> relRoots = sqlConverter.sqlToRelNode(sqlQuery, catalogReader);
183-
assertEquals(1, relRoots.size());
184-
RelRoot calcite1 = relRoots.get(0);
181+
// SQL -> Substrait Plan 1
182+
Plan plan1 = sqlConverter.convert(sqlQuery, catalogReader);
183+
assertEquals(1, plan1.getRoots().size());
185184

186-
// Calcite 1 -> Substrait POJO 1
187-
Plan.Root pojo1 = SubstraitRelVisitor.convert(calcite1, extensions);
185+
// Substrait Plan 1 -> Substrait Root 1
186+
Plan.Root root1 = plan1.getRoots().get(0);
188187

189-
// Substrait POJO 1 -> Substrait Proto
190-
io.substrait.proto.RelRoot proto = new RelProtoConverter(extensionCollector).toProto(pojo1);
188+
// Substrait Root 1 -> Substrait Proto
189+
io.substrait.proto.RelRoot proto = new RelProtoConverter(extensionCollector).toProto(root1);
191190

192-
// Substrait Proto -> Substrait Pojo 2
193-
Plan.Root pojo2 = new ProtoRelConverter(extensionCollector, extensions).from(proto);
191+
// Substrait Proto -> Substrait Root 2
192+
Plan.Root root2 = new ProtoRelConverter(extensionCollector, extensions).from(proto);
194193

195-
// Verify that POJOs are the same
196-
assertEquals(pojo1, pojo2);
194+
// Verify that roots are the same
195+
assertEquals(root1, root2);
197196

198-
// Substrait POJO 2 -> Calcite 2
197+
// Substrait Root 2 -> Calcite 2
199198
final SubstraitToCalcite substraitToCalcite =
200199
new SubstraitToCalcite(extensions, typeFactory, catalogReader);
201200

202-
RelRoot calcite2 = substraitToCalcite.convert(pojo2);
201+
RelRoot calcite2 = substraitToCalcite.convert(root2);
203202
// It would be ideal to compare calcite1 and calcite2, however there isn't a good mechanism to
204203
// do so
205204
assertNotNull(calcite2);
206205

207-
// Calcite 2 -> Substrait POJO 3
208-
Plan.Root pojo3 = SubstraitRelVisitor.convert(calcite2, extensions);
206+
// Calcite 2 -> Substrait Root 3
207+
Plan.Root root3 = SubstraitRelVisitor.convert(calcite2, extensions);
209208

210209
// Verify that POJOs are the same
211-
assertEquals(pojo1, pojo3);
210+
assertEquals(root1, root3);
212211
}
213212

214213
/**
@@ -285,9 +284,9 @@ protected void assertRowMatch(RelDataType actual, List<Type> expected) {
285284
assertEquals(expected, struct.fields());
286285
}
287286

288-
protected io.substrait.proto.Plan toSubstraitPlan(String sql, CalciteCatalogReader catalog)
287+
protected Plan toSubstraitPlan(String sql, CalciteCatalogReader catalog)
289288
throws SqlParseException {
290-
return new SqlToSubstrait().execute(sql, catalog);
289+
return new SqlToSubstrait().convert(sql, catalog);
291290
}
292291

293292
protected String toSql(io.substrait.proto.Plan protoPlan) {
@@ -305,6 +304,10 @@ protected String toSql(Plan plan) {
305304
return SubstraitSqlDialect.toSql(project).getSql();
306305
}
307306

307+
protected io.substrait.proto.Plan toProto(Plan plan) {
308+
return new PlanProtoConverter().toProto(plan);
309+
}
310+
308311
protected static CalciteCatalogReader schemaToCatalog(String schemaName, Schema schema) {
309312
CalciteSchema rootSchema = CalciteSchema.createRootSchema(false);
310313
rootSchema.add(schemaName, schema);

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
import io.substrait.isthmus.utils.SetUtils;
66
import io.substrait.plan.Plan;
7-
import io.substrait.plan.PlanProtoConverter;
87
import io.substrait.plan.ProtoPlanConverter;
98
import io.substrait.proto.AggregateFunction;
109
import io.substrait.relation.Cross;
@@ -23,7 +22,7 @@ public class ProtoPlanConverterTest extends PlanTestBase {
2322

2423
private io.substrait.proto.Plan getProtoPlan(String query1) throws SqlParseException {
2524
SqlToSubstrait s = new SqlToSubstrait();
26-
return s.execute(query1, TPCH_CATALOG);
25+
return toProto(s.convert(query1, TPCH_CATALOG));
2726
}
2827

2928
@Test
@@ -54,8 +53,7 @@ public void distinctCount() throws IOException, SqlParseException {
5453
String distinctQuery = "select count(DISTINCT L_ORDERKEY) from lineitem";
5554
io.substrait.proto.Plan protoPlan = getProtoPlan(distinctQuery);
5655
assertAggregateInvocationDistinct(protoPlan);
57-
assertAggregateInvocationDistinct(
58-
new PlanProtoConverter().toProto(new ProtoPlanConverter().from(protoPlan)));
56+
assertAggregateInvocationDistinct(toProto(new ProtoPlanConverter().from(protoPlan)));
5957
}
6058

6159
@Test

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
import io.substrait.extension.SimpleExtension;
1010
import io.substrait.isthmus.sql.SubstraitSqlDialect;
1111
import io.substrait.plan.Plan;
12-
import io.substrait.plan.ProtoPlanConverter;
1312
import io.substrait.relation.Aggregate;
1413
import io.substrait.relation.CopyOnWriteUtils;
1514
import io.substrait.relation.NamedScan;
@@ -77,8 +76,7 @@ public class RelCopyOnWriteVisitorTest extends PlanTestBase {
7776

7877
private Plan buildPlanFromQuery(String query) throws IOException, SqlParseException {
7978
SqlToSubstrait s = new SqlToSubstrait();
80-
io.substrait.proto.Plan protoPlan1 = s.execute(query, TPCH_CATALOG);
81-
return new ProtoPlanConverter().from(protoPlan1);
79+
return s.convert(query, TPCH_CATALOG);
8280
}
8381

8482
@Test

0 commit comments

Comments
 (0)