Skip to content

Commit 8cd5350

Browse files
ZorinAntonZorinAnton
committed
chore(isthmus): rebase + refactoring
1 parent 135c6ff commit 8cd5350

8 files changed

Lines changed: 333 additions & 298 deletions

File tree

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

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
package io.substrait.isthmus;
2+
3+
import io.substrait.isthmus.operation.CalciteOperation;
4+
import io.substrait.isthmus.operation.CalciteOperationBuilder;
5+
import io.substrait.isthmus.operation.RelationalOperation;
6+
import io.substrait.isthmus.sql.SubstraitSqlStatementParser;
7+
import io.substrait.isthmus.sql.SubstraitSqlValidator;
8+
import java.util.List;
9+
import java.util.stream.Collectors;
10+
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
11+
import org.apache.calcite.plan.RelOptCluster;
12+
import org.apache.calcite.plan.RelOptPlanner;
13+
import org.apache.calcite.plan.hep.HepPlanner;
14+
import org.apache.calcite.plan.hep.HepProgram;
15+
import org.apache.calcite.prepare.Prepare;
16+
import org.apache.calcite.rel.RelNode;
17+
import org.apache.calcite.rel.RelRoot;
18+
import org.apache.calcite.rel.rules.CoreRules;
19+
import org.apache.calcite.rex.RexBuilder;
20+
import org.apache.calcite.sql.SqlNode;
21+
import org.apache.calcite.sql.parser.SqlParseException;
22+
import org.apache.calcite.sql.parser.SqlParser;
23+
import org.apache.calcite.sql.parser.ddl.SqlDdlParserImpl;
24+
import org.apache.calcite.sql.validate.SqlConformanceEnum;
25+
import org.apache.calcite.sql.validate.SqlValidator;
26+
import org.apache.calcite.sql2rel.SqlToRelConverter;
27+
import org.apache.calcite.sql2rel.StandardConvertletTable;
28+
29+
public class SqlToCalcite {
30+
protected static FeatureBoard createDefaultFeatureBoard() {
31+
return ImmutableFeatureBoard.builder().build();
32+
}
33+
34+
protected static SqlParser.Config createDefaultParserConfig() {
35+
return SqlParser.Config.DEFAULT
36+
.withUnquotedCasing(createDefaultFeatureBoard().unquotedCasing())
37+
.withParserFactory(SqlDdlParserImpl.FACTORY)
38+
.withConformance(SqlConformanceEnum.LENIENT);
39+
}
40+
41+
protected static SqlParser.Config createDefaultParserConfig(FeatureBoard featureBoard) {
42+
return SqlParser.Config.DEFAULT
43+
.withUnquotedCasing(featureBoard.unquotedCasing())
44+
.withParserFactory(SqlDdlParserImpl.FACTORY)
45+
.withConformance(SqlConformanceEnum.LENIENT);
46+
}
47+
48+
protected static RelOptCluster createDefaultRelOptCluster() {
49+
RexBuilder rexBuilder =
50+
new RexBuilder(new JavaTypeFactoryImpl(SubstraitTypeSystem.TYPE_SYSTEM));
51+
HepProgram program = HepProgram.builder().build();
52+
RelOptPlanner emptyPlanner = new HepPlanner(program);
53+
return RelOptCluster.create(emptyPlanner, rexBuilder);
54+
}
55+
56+
protected static SqlToRelConverter createSqlToRelConverter(
57+
SqlValidator validator, Prepare.CatalogReader catalogReader, RelOptCluster relOptCluster) {
58+
return new SqlToRelConverter(
59+
null,
60+
validator,
61+
catalogReader,
62+
relOptCluster,
63+
StandardConvertletTable.INSTANCE,
64+
SqlToRelConverter.CONFIG);
65+
}
66+
67+
protected static RelRoot parseRelationalExpression(SqlToRelConverter converter, SqlNode sqlNode) {
68+
RelRoot converted = converter.convertQuery(sqlNode, true, true);
69+
return converted.withRel(removeRedundantProjects(converted.rel));
70+
}
71+
72+
protected static RelNode removeRedundantProjects(RelNode root) {
73+
// The Calcite RelBuilder, when constructing Project that does not modify its inputs in any way,
74+
// simply elides it. The PROJECT_REMOVE rule can be used to remove such projects from Rel trees.
75+
// This facilitates roundtrip testing.
76+
HepProgram program = HepProgram.builder().addRuleInstance(CoreRules.PROJECT_REMOVE).build();
77+
HepPlanner planner = new HepPlanner(program);
78+
planner.setRoot(root);
79+
return planner.findBestExp();
80+
}
81+
82+
protected static CalciteOperation convert(
83+
SqlNode sqlNode,
84+
Prepare.CatalogReader catalogReader,
85+
SqlValidator validator,
86+
RelOptCluster cluster) {
87+
SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader, cluster);
88+
CalciteOperationBuilder operationBuilder =
89+
new CalciteOperationBuilder(converter, SqlToCalcite::parseRelationalExpression);
90+
return sqlNode.accept(operationBuilder);
91+
}
92+
93+
public static CalciteOperation convertQuery(
94+
String sqlStatement, Prepare.CatalogReader catalogReader) throws SqlParseException {
95+
return convertQuery(sqlStatement, catalogReader, createDefaultParserConfig());
96+
}
97+
98+
public static CalciteOperation convertQuery(
99+
String sqlStatement, Prepare.CatalogReader catalogReader, SqlParser.Config parserConfig)
100+
throws SqlParseException {
101+
List<SqlNode> sqlNodes =
102+
SubstraitSqlStatementParser.parseStatements(sqlStatement, parserConfig);
103+
if (sqlNodes.size() != 1) {
104+
throw new IllegalArgumentException(
105+
String.format("Expected one statement, found %d", sqlNodes.size()));
106+
}
107+
SqlValidator validator = new SubstraitSqlValidator(catalogReader);
108+
return convertQuery(sqlNodes.get(0), catalogReader, validator, createDefaultRelOptCluster());
109+
}
110+
111+
public static RelRoot convertRelationalQuery(
112+
String sqlStatement, Prepare.CatalogReader catalogReader) throws SqlParseException {
113+
return convertRelationalQuery(sqlStatement, catalogReader, createDefaultParserConfig());
114+
}
115+
116+
public static RelRoot convertRelationalQuery(
117+
String sqlStatement, Prepare.CatalogReader catalogReader, SqlParser.Config parserConfig)
118+
throws SqlParseException {
119+
CalciteOperation calciteOperation = convertQuery(sqlStatement, catalogReader, parserConfig);
120+
if (calciteOperation instanceof RelationalOperation) {
121+
return ((RelationalOperation) calciteOperation).getRelRoot();
122+
}
123+
throw new IllegalArgumentException(
124+
String.format("Non-relational algebra statement: %s", sqlStatement));
125+
}
126+
127+
public static CalciteOperation convertQuery(
128+
SqlNode sqlNode,
129+
Prepare.CatalogReader catalogReader,
130+
SqlValidator validator,
131+
RelOptCluster cluster) {
132+
return convert(sqlNode, catalogReader, validator, cluster);
133+
}
134+
135+
public static List<CalciteOperation> convertQueries(
136+
String sqlStatements, Prepare.CatalogReader catalogReader) throws SqlParseException {
137+
List<SqlNode> sqlNodes =
138+
SubstraitSqlStatementParser.parseStatements(sqlStatements, createDefaultParserConfig());
139+
SqlValidator validator = new SubstraitSqlValidator(catalogReader);
140+
RelOptCluster cluster = createDefaultRelOptCluster();
141+
142+
return sqlNodes.stream()
143+
.map(sqlNode -> convert(sqlNode, catalogReader, validator, cluster))
144+
.collect(Collectors.toUnmodifiableList());
145+
}
146+
147+
public static List<RelRoot> convertRelationalQueries(
148+
String sqlStatements, Prepare.CatalogReader catalogReader) throws SqlParseException {
149+
return convertQueries(sqlStatements, catalogReader).stream()
150+
.filter(calciteOperation -> calciteOperation instanceof RelationalOperation)
151+
.map(calciteOperation -> ((RelationalOperation) calciteOperation).getRelRoot())
152+
.collect(Collectors.toUnmodifiableList());
153+
}
154+
}

0 commit comments

Comments
 (0)