Skip to content

Commit fe86d02

Browse files
ZorinAntonZorinAnton
committed
a first draft
1 parent 1f483dd commit fe86d02

File tree

3 files changed

+194
-9
lines changed

3 files changed

+194
-9
lines changed

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

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

33
import com.google.common.annotations.VisibleForTesting;
4+
import io.substrait.isthmus.expression.DdlRelBuilder;
5+
import io.substrait.plan.ImmutablePlan;
46
import io.substrait.plan.Plan.Version;
57
import io.substrait.plan.PlanProtoConverter;
68
import io.substrait.proto.Plan;
9+
import java.util.LinkedList;
710
import java.util.List;
811
import org.apache.calcite.plan.hep.HepPlanner;
912
import org.apache.calcite.plan.hep.HepProgram;
@@ -12,6 +15,7 @@
1215
import org.apache.calcite.rel.RelRoot;
1316
import org.apache.calcite.schema.Schema;
1417
import org.apache.calcite.sql.SqlNode;
18+
import org.apache.calcite.sql.SqlNodeList;
1519
import org.apache.calcite.sql.parser.SqlParseException;
1620
import org.apache.calcite.sql.parser.SqlParser;
1721
import org.apache.calcite.sql.validate.SqlValidator;
@@ -59,26 +63,64 @@ private Plan executeInner(String sql, SqlValidator validator, Prepare.CatalogRea
5963
builder.version(Version.builder().from(Version.DEFAULT_VERSION).producer("isthmus").build());
6064

6165
// TODO: consider case in which one sql passes conversion while others don't
62-
sqlToRelNode(sql, validator, catalogReader).stream()
63-
.map(root -> SubstraitRelVisitor.convert(root, EXTENSION_COLLECTION, featureBoard))
64-
.forEach(root -> builder.addRoots(root));
65-
66+
sqlToRootNodes(sql, validator, catalogReader, builder);
6667
PlanProtoConverter planToProto = new PlanProtoConverter();
6768

6869
return planToProto.toProto(builder.build());
6970
}
7071

72+
private void sqlToRootNodes(
73+
String sql,
74+
SqlValidator validator,
75+
Prepare.CatalogReader catalogReader,
76+
ImmutablePlan.Builder builder)
77+
throws SqlParseException {
78+
SqlParser parser = SqlParser.create(sql, parserConfig);
79+
var parsedList = parser.parseStmtList();
80+
SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);
81+
// IMPORTANT: parsedList gets filtered in the call below
82+
List<io.substrait.plan.Plan.Root> ddlRelRoot = ddlSqlToRootNodes(parsedList, converter);
83+
ddlRelRoot.forEach(builder::addRoots);
84+
85+
sqlNodesToRelNode(parsedList, converter).stream()
86+
.map(relRoot -> SubstraitRelVisitor.convert(relRoot, EXTENSION_COLLECTION, featureBoard))
87+
.forEach(builder::addRoots);
88+
}
89+
90+
private List<io.substrait.plan.Plan.Root> ddlSqlToRootNodes(
91+
final SqlNodeList sqlNodeList, final SqlToRelConverter converter) throws SqlParseException {
92+
93+
final DdlRelBuilder ddlRelBuilder =
94+
new DdlRelBuilder(
95+
converter, SqlToSubstrait::getBestExpRelRoot, EXTENSION_COLLECTION, featureBoard);
96+
97+
List<SqlNode> toRemove = new LinkedList<>();
98+
List<io.substrait.plan.Plan.Root> retVal = new LinkedList<>();
99+
for (final SqlNode sqlNode : sqlNodeList) {
100+
final io.substrait.plan.Plan.Root root = sqlNode.accept(ddlRelBuilder);
101+
if (root != null) {
102+
retVal.add(root);
103+
toRemove.add(sqlNode);
104+
}
105+
}
106+
sqlNodeList.removeAll(toRemove);
107+
return retVal;
108+
}
109+
110+
private List<RelRoot> sqlNodesToRelNode(
111+
final SqlNodeList parsedList, final SqlToRelConverter converter) {
112+
return parsedList.stream()
113+
.map(parsed -> getBestExpRelRoot(converter, parsed))
114+
.collect(java.util.stream.Collectors.toList());
115+
}
116+
71117
private List<RelRoot> sqlToRelNode(
72118
String sql, SqlValidator validator, Prepare.CatalogReader catalogReader)
73119
throws SqlParseException {
74120
SqlParser parser = SqlParser.create(sql, parserConfig);
75121
var parsedList = parser.parseStmtList();
76122
SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);
77-
List<RelRoot> roots =
78-
parsedList.stream()
79-
.map(parsed -> getBestExpRelRoot(converter, parsed))
80-
.collect(java.util.stream.Collectors.toList());
81-
return roots;
123+
return sqlNodesToRelNode(parsedList, converter);
82124
}
83125

84126
@VisibleForTesting

‎isthmus/src/main/java/io/substrait/isthmus/expression/DdlRelBuilder.java‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package io.substrait.isthmus.expression;
2+
3+
import io.substrait.expression.Expression;
4+
import io.substrait.expression.ExpressionCreator;
5+
import io.substrait.extension.SimpleExtension;
6+
import io.substrait.isthmus.FeatureBoard;
7+
import io.substrait.isthmus.SubstraitRelVisitor;
8+
import io.substrait.isthmus.TypeConverter;
9+
import io.substrait.plan.Plan;
10+
import io.substrait.relation.AbstractDdlRel;
11+
import io.substrait.relation.AbstractWriteRel;
12+
import io.substrait.relation.NamedDdl;
13+
import io.substrait.relation.NamedWrite;
14+
import io.substrait.type.NamedStruct;
15+
import java.util.Map;
16+
import java.util.concurrent.ConcurrentHashMap;
17+
import java.util.function.BiFunction;
18+
import java.util.function.Function;
19+
import org.apache.calcite.rel.RelRoot;
20+
import org.apache.calcite.rel.type.RelDataType;
21+
import org.apache.calcite.sql.SqlCall;
22+
import org.apache.calcite.sql.SqlNode;
23+
import org.apache.calcite.sql.ddl.SqlCreateTable;
24+
import org.apache.calcite.sql.ddl.SqlCreateView;
25+
import org.apache.calcite.sql.util.SqlBasicVisitor;
26+
import org.apache.calcite.sql2rel.SqlToRelConverter;
27+
28+
public class DdlRelBuilder extends SqlBasicVisitor<Plan.Root> {
29+
protected final Map<Class<? extends SqlCall>, Function<SqlCall, Plan.Root>> createHandlers =
30+
new ConcurrentHashMap<>();
31+
32+
private final SqlToRelConverter converter;
33+
private final BiFunction<SqlToRelConverter, SqlNode, RelRoot> bestExpRelRootGettter;
34+
private final SimpleExtension.ExtensionCollection extensionCollection;
35+
private final FeatureBoard featureBoard;
36+
37+
private Function<SqlCall, Plan.Root> findCreateHandler(final SqlCall call) {
38+
Class<?> currentClass = call.getClass();
39+
while (SqlCall.class.isAssignableFrom(currentClass)) {
40+
final Function<SqlCall, Plan.Root> found = createHandlers.get(currentClass);
41+
if (found != null) {
42+
return found;
43+
}
44+
currentClass = currentClass.getSuperclass();
45+
}
46+
return null;
47+
}
48+
49+
public DdlRelBuilder(
50+
final SqlToRelConverter converter,
51+
final BiFunction<SqlToRelConverter, SqlNode, RelRoot> bestExpRelRootGetter,
52+
final SimpleExtension.ExtensionCollection extensionCollection,
53+
final FeatureBoard featureBoard) {
54+
super();
55+
this.converter = converter;
56+
this.bestExpRelRootGettter = bestExpRelRootGetter;
57+
this.extensionCollection = extensionCollection;
58+
this.featureBoard = featureBoard;
59+
60+
createHandlers.put(
61+
SqlCreateTable.class, sqlCall -> handleCreateTable((SqlCreateTable) sqlCall));
62+
createHandlers.put(SqlCreateView.class, sqlCall -> handleCreateView((SqlCreateView) sqlCall));
63+
}
64+
65+
@Override
66+
public Plan.Root visit(final SqlCall sqlCall) {
67+
Function<SqlCall, Plan.Root> createHandler = findCreateHandler(sqlCall);
68+
if (createHandler == null) {
69+
return null;
70+
}
71+
72+
return createHandler.apply(sqlCall);
73+
}
74+
75+
private NamedStruct getSchema(final RelRoot queryRelRoot) {
76+
final RelDataType rowType = queryRelRoot.rel.getRowType();
77+
78+
final TypeConverter typeConverter = TypeConverter.DEFAULT;
79+
return typeConverter.toNamedStruct(rowType);
80+
}
81+
82+
public Plan.Root handleCreateTable(final SqlCreateTable sqlCreateTable) {
83+
if (sqlCreateTable.query == null) {
84+
throw new IllegalArgumentException("Only create table as select statements are supported");
85+
}
86+
87+
final RelRoot queryRelRoot = bestExpRelRootGettter.apply(converter, sqlCreateTable.query);
88+
89+
NamedStruct schema = getSchema(queryRelRoot);
90+
91+
var rel = SubstraitRelVisitor.convert(queryRelRoot, extensionCollection, featureBoard);
92+
NamedWrite namedWrite =
93+
NamedWrite.builder()
94+
.input(rel.getInput())
95+
.tableSchema(schema)
96+
.operation(AbstractWriteRel.WriteOp.CTAS)
97+
.createMode(AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS)
98+
.outputMode(AbstractWriteRel.OutputMode.NO_OUTPUT)
99+
.names(sqlCreateTable.name.names)
100+
.build();
101+
102+
return Plan.Root.builder().input(namedWrite).build();
103+
}
104+
105+
Plan.Root handleCreateView(final SqlCreateView sqlCreateView) {
106+
107+
final RelRoot queryRelRoot = bestExpRelRootGettter.apply(converter, sqlCreateView.query);
108+
var rel = SubstraitRelVisitor.convert(queryRelRoot, extensionCollection, featureBoard);
109+
final Expression.StructLiteral defaults = ExpressionCreator.struct(true);
110+
111+
final NamedDdl namedDdl =
112+
NamedDdl.builder()
113+
.viewDefinition(rel.getInput())
114+
.tableSchema(getSchema(queryRelRoot))
115+
.tableDefaults(defaults)
116+
.operation(AbstractDdlRel.DdlOp.CREATE)
117+
.object(AbstractDdlRel.DdlObject.VIEW)
118+
.names(sqlCreateView.name.names)
119+
.build();
120+
121+
return Plan.Root.builder().input(namedDdl).build();
122+
}
123+
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package io.substrait.isthmus;
2+
3+
import io.substrait.proto.Plan;
4+
import java.util.List;
5+
import org.apache.calcite.sql.parser.SqlParseException;
6+
import org.junit.jupiter.api.Test;
7+
8+
public class SqlToSubstraitTest {
9+
10+
@Test
11+
void test() throws SqlParseException {
12+
SqlToSubstrait s = new SqlToSubstrait();
13+
Plan plan =
14+
s.execute(
15+
"create view test1 as select * from test; "
16+
+ "create table test2 as select * from test; ",
17+
List.of("create table test (intcol int, charcol varchar(10))"));
18+
System.out.println(plan.toString());
19+
}
20+
}

0 commit comments

Comments
 (0)