Skip to content

Commit 1a0cbf2

Browse files
ZorinAntonZorinAnton
committed
refactor(isthmus): introduce CacliteOperation to represent rel algebra and ddl
1 parent 2518d55 commit 1a0cbf2

16 files changed

Lines changed: 745 additions & 11 deletions

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

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
package io.substrait.isthmus;
2+
3+
import io.substrait.expression.Expression;
4+
import io.substrait.expression.ExpressionCreator;
5+
import io.substrait.extension.SimpleExtension;
6+
import io.substrait.isthmus.operation.CalciteOperationBasicVisitor;
7+
import io.substrait.isthmus.operation.CreateTableAs;
8+
import io.substrait.isthmus.operation.CreateView;
9+
import io.substrait.isthmus.operation.RelationalOperation;
10+
import io.substrait.plan.Plan;
11+
import io.substrait.relation.AbstractDdlRel;
12+
import io.substrait.relation.AbstractWriteRel;
13+
import io.substrait.relation.NamedDdl;
14+
import io.substrait.relation.NamedWrite;
15+
import io.substrait.type.NamedStruct;
16+
import org.apache.calcite.rel.RelRoot;
17+
import org.apache.calcite.rel.type.RelDataType;
18+
19+
public class CalciteOperationToSubstrait extends CalciteOperationBasicVisitor<Plan.Root> {
20+
private final SimpleExtension.ExtensionCollection extensionCollection;
21+
private final FeatureBoard featureBoard;
22+
23+
public CalciteOperationToSubstrait(
24+
SimpleExtension.ExtensionCollection extensionCollection, FeatureBoard featureBoard) {
25+
this.extensionCollection = extensionCollection;
26+
this.featureBoard = featureBoard;
27+
}
28+
29+
private NamedStruct getSchema(final RelRoot queryRelRoot) {
30+
final RelDataType rowType = queryRelRoot.rel.getRowType();
31+
32+
final TypeConverter typeConverter = TypeConverter.DEFAULT;
33+
return typeConverter.toNamedStruct(rowType);
34+
}
35+
36+
@Override
37+
public Plan.Root visit(RelationalOperation relationalOperation) {
38+
return SubstraitRelVisitor.convert(
39+
relationalOperation.getRelRoot(), extensionCollection, featureBoard);
40+
}
41+
42+
@Override
43+
public Plan.Root visit(CreateTableAs createTableAs) {
44+
RelRoot input = createTableAs.getInput();
45+
Plan.Root rel = SubstraitRelVisitor.convert(input, extensionCollection, featureBoard);
46+
NamedStruct schema = getSchema(input);
47+
48+
NamedWrite namedWrite =
49+
NamedWrite.builder()
50+
.input(rel.getInput())
51+
.tableSchema(schema)
52+
.operation(AbstractWriteRel.WriteOp.CTAS)
53+
.createMode(AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS)
54+
.outputMode(AbstractWriteRel.OutputMode.NO_OUTPUT)
55+
.names(createTableAs.getNames())
56+
.build();
57+
58+
return Plan.Root.builder().input(namedWrite).build();
59+
}
60+
61+
@Override
62+
public Plan.Root visit(CreateView createView) {
63+
RelRoot input = createView.getInput();
64+
Plan.Root rel = SubstraitRelVisitor.convert(input, extensionCollection, featureBoard);
65+
final Expression.StructLiteral defaults = ExpressionCreator.struct(false);
66+
67+
final NamedDdl namedDdl =
68+
NamedDdl.builder()
69+
.viewDefinition(rel.getInput())
70+
.tableSchema(getSchema(input))
71+
.tableDefaults(defaults)
72+
.operation(AbstractDdlRel.DdlOp.CREATE)
73+
.object(AbstractDdlRel.DdlObject.VIEW)
74+
.names(createView.getNames())
75+
.build();
76+
77+
return Plan.Root.builder().input(namedDdl).build();
78+
}
79+
}

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import io.substrait.plan.Plan;
88
import io.substrait.plan.Plan.Version;
99
import io.substrait.plan.PlanProtoConverter;
10-
import java.util.ArrayList;
1110
import java.util.List;
1211
import org.apache.calcite.plan.hep.HepPlanner;
1312
import org.apache.calcite.plan.hep.HepProgram;
@@ -110,7 +109,18 @@ List<RelRoot> sqlToRelNode(String sql, Prepare.CatalogReader catalogReader)
110109
return roots;
111110
}
112111

113-
protected void sqlToPlanRoots(
112+
void sqlToPlanRoots(
113+
String sql, SqlValidator validator, Prepare.CatalogReader catalogReader, Builder builder)
114+
throws SqlParseException {
115+
List<CalciteOperation> calciteOperations = sqlToCalciteOperation(sql, validator, catalogReader);
116+
CalciteOperationToSubstrait calciteOperationToSubstrait =
117+
new CalciteOperationToSubstrait(EXTENSION_COLLECTION, featureBoard);
118+
calciteOperations.stream()
119+
.map(operation -> operation.accept(calciteOperationToSubstrait))
120+
.forEach(builder::addRoots);
121+
}
122+
123+
void sqlToPlanRoots1(
114124
String sql, SqlValidator validator, Prepare.CatalogReader catalogReader, Builder builder)
115125
throws SqlParseException {
116126

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

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -588,11 +588,8 @@ public RelNode visit(VirtualTableScan virtualTableScan, Context context) {
588588
public RelNode visit(NamedWrite write, Context context) {
589589
RelNode input = write.getInput().accept(this, context);
590590
assert relBuilder.getRelOptSchema() != null;
591-
final RelOptTable table = relBuilder.getRelOptSchema().getTableForMember(write.getNames());
592-
593-
if (table == null) {
594-
throw new IllegalStateException("Table not found in Calcite catalog: " + write.getNames());
595-
}
591+
final RelOptTable targetTable =
592+
relBuilder.getRelOptSchema().getTableForMember(write.getNames());
596593

597594
TableModify.Operation operation;
598595
switch (write.getOperation()) {
@@ -610,8 +607,12 @@ public RelNode visit(NamedWrite write, Context context) {
610607
+ "Check if a more specific relation type (e.g., NamedUpdate) should be used.");
611608
}
612609

610+
if (targetTable == null) {
611+
throw new IllegalStateException("Table not found in Calcite catalog: " + write.getNames());
612+
}
613+
613614
return LogicalTableModify.create(
614-
table,
615+
targetTable,
615616
(Prepare.CatalogReader) relBuilder.getRelOptSchema(),
616617
input,
617618
operation,

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package io.substrait.isthmus;
2+
3+
import io.substrait.extension.SimpleExtension;
4+
import io.substrait.isthmus.operation.CalciteOperation;
5+
import io.substrait.isthmus.operation.CreateTableAs;
6+
import io.substrait.isthmus.operation.CreateView;
7+
import io.substrait.isthmus.operation.RelationalOperation;
8+
import io.substrait.isthmus.operation.SqlKindFromRel;
9+
import io.substrait.relation.AbstractRelVisitor;
10+
import io.substrait.relation.NamedDdl;
11+
import io.substrait.relation.NamedWrite;
12+
import io.substrait.relation.Rel;
13+
import io.substrait.util.EmptyVisitationContext;
14+
import java.util.List;
15+
import org.apache.calcite.plan.RelTraitDef;
16+
import org.apache.calcite.prepare.Prepare;
17+
import org.apache.calcite.rel.RelNode;
18+
import org.apache.calcite.rel.RelRoot;
19+
import org.apache.calcite.rel.type.RelDataTypeFactory;
20+
import org.apache.calcite.sql.SqlKind;
21+
import org.apache.calcite.sql.parser.SqlParser;
22+
import org.apache.calcite.tools.Frameworks;
23+
import org.apache.calcite.tools.RelBuilder;
24+
25+
public class SubstraitToCalciteOperation
26+
extends AbstractRelVisitor<
27+
CalciteOperation, SubstraitRelNodeConverter.Context, RuntimeException> {
28+
29+
private final SubstraitRelNodeConverter substraitRelNodeConverter;
30+
31+
public SubstraitToCalciteOperation(
32+
SimpleExtension.ExtensionCollection extensionCollection,
33+
RelDataTypeFactory relDataTypeFactory,
34+
Prepare.CatalogReader catalogReader,
35+
SqlParser.Config parserConfig) {
36+
37+
RelBuilder relBuilder =
38+
RelBuilder.create(
39+
Frameworks.newConfigBuilder()
40+
.parserConfig(parserConfig)
41+
.defaultSchema(catalogReader.getRootSchema().plus())
42+
.traitDefs((List<RelTraitDef>) null)
43+
.programs()
44+
.build());
45+
this.substraitRelNodeConverter =
46+
new SubstraitRelNodeConverter(extensionCollection, relDataTypeFactory, relBuilder);
47+
}
48+
49+
@Override
50+
public CalciteOperation visit(NamedDdl namedDdl, SubstraitRelNodeConverter.Context ctx) {
51+
if (namedDdl.getViewDefinition().isEmpty()) {
52+
throw new IllegalArgumentException("no view definition found");
53+
}
54+
Rel viewDefinition = namedDdl.getViewDefinition().get();
55+
RelNode relNode = viewDefinition.accept(substraitRelNodeConverter, ctx);
56+
RelRoot relRoot = RelRoot.of(relNode, SqlKind.CREATE_VIEW);
57+
58+
return new CreateView(namedDdl.getNames(), relRoot);
59+
}
60+
61+
@Override
62+
public CalciteOperation visit(NamedWrite namedWrite, SubstraitRelNodeConverter.Context ctx) {
63+
Rel input = namedWrite.getInput();
64+
RelNode relNode = input.accept(substraitRelNodeConverter, ctx);
65+
RelRoot relRoot = RelRoot.of(relNode, SqlKind.CREATE_TABLE);
66+
return new CreateTableAs(namedWrite.getNames(), relRoot);
67+
}
68+
69+
@Override
70+
public CalciteOperation visitFallback(Rel rel, SubstraitRelNodeConverter.Context context) {
71+
SqlKindFromRel sqlKindFromRel = new SqlKindFromRel();
72+
SqlKind kind = rel.accept(sqlKindFromRel, EmptyVisitationContext.INSTANCE);
73+
RelNode relNode = rel.accept(substraitRelNodeConverter, context);
74+
RelRoot relRoot = RelRoot.of(relNode, kind);
75+
return new RelationalOperation(relRoot);
76+
}
77+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/CalciteOperation.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package io.substrait.isthmus.operation;
2+
3+
public class CalciteOperation {
4+
public <R> R accept(CalciteOperationVisitor<R> visitor) {
5+
return visitor.visit(this);
6+
}
7+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/CalciteOperationBasicVisitor.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package io.substrait.isthmus.operation;
2+
3+
public class CalciteOperationBasicVisitor<R> implements CalciteOperationVisitor<R> {
4+
@Override
5+
public R visit(CalciteOperation operation) {
6+
return null;
7+
}
8+
9+
@Override
10+
public R visit(RelationalOperation relationalOperation) {
11+
return visit((CalciteOperation) relationalOperation);
12+
}
13+
14+
@Override
15+
public R visit(DdlOperation ddlOperation) {
16+
return visit((CalciteOperation) ddlOperation);
17+
}
18+
19+
@Override
20+
public R visit(CreateWithInput createWithInput) {
21+
return visit((Create) createWithInput);
22+
}
23+
24+
@Override
25+
public R visit(Create create) {
26+
return visit((DdlOperation) create);
27+
}
28+
29+
@Override
30+
public R visit(CreateTableAs createTableAs) {
31+
return visit((CreateWithInput) createTableAs);
32+
}
33+
34+
@Override
35+
public R visit(CreateView createView) {
36+
return visit((CreateWithInput) createView);
37+
}
38+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/CalciteOperationBuilder.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package io.substrait.isthmus.operation;
2+
3+
import java.util.Map;
4+
import java.util.concurrent.ConcurrentHashMap;
5+
import java.util.function.BiFunction;
6+
import java.util.function.Function;
7+
import org.apache.calcite.rel.RelRoot;
8+
import org.apache.calcite.sql.SqlCall;
9+
import org.apache.calcite.sql.SqlNode;
10+
import org.apache.calcite.sql.ddl.SqlCreateTable;
11+
import org.apache.calcite.sql.ddl.SqlCreateView;
12+
import org.apache.calcite.sql.util.SqlBasicVisitor;
13+
import org.apache.calcite.sql2rel.SqlToRelConverter;
14+
15+
public class CalciteOperationBuilder extends SqlBasicVisitor<CalciteOperation> {
16+
protected final Map<Class<? extends SqlCall>, Function<SqlCall, CalciteOperation>> ddlHandlers =
17+
new ConcurrentHashMap<>();
18+
private final SqlToRelConverter converter;
19+
private final BiFunction<SqlToRelConverter, SqlNode, RelRoot> bestExprRelRootBuilder;
20+
21+
public CalciteOperationBuilder(
22+
final SqlToRelConverter converter,
23+
BiFunction<SqlToRelConverter, SqlNode, RelRoot> bestExprRelRootBuilder) {
24+
this.converter = converter;
25+
this.bestExprRelRootBuilder = bestExprRelRootBuilder;
26+
27+
ddlHandlers.put(SqlCreateTable.class, sqlCall -> handleCreateTable((SqlCreateTable) sqlCall));
28+
ddlHandlers.put(SqlCreateView.class, sqlCall -> handleCreateView((SqlCreateView) sqlCall));
29+
}
30+
31+
private Function<SqlCall, CalciteOperation> findDdlHandler(final SqlCall call) {
32+
Class<?> currentClass = call.getClass();
33+
while (SqlCall.class.isAssignableFrom(currentClass)) {
34+
final Function<SqlCall, CalciteOperation> found = ddlHandlers.get(currentClass);
35+
if (found != null) {
36+
return found;
37+
}
38+
currentClass = currentClass.getSuperclass();
39+
}
40+
return null;
41+
}
42+
43+
@Override
44+
public CalciteOperation visit(final SqlCall sqlCall) {
45+
Function<SqlCall, CalciteOperation> ddlHandler = findDdlHandler(sqlCall);
46+
if (ddlHandler != null) {
47+
return ddlHandler.apply(sqlCall);
48+
}
49+
return handleRelationalOperation(sqlCall);
50+
}
51+
52+
protected CalciteOperation handleRelationalOperation(final SqlNode sqlNode) {
53+
return new RelationalOperation(bestExprRelRootBuilder.apply(converter, sqlNode));
54+
}
55+
56+
protected CalciteOperation handleCreateTable(final SqlCreateTable sqlCreateTable) {
57+
if (sqlCreateTable.query == null) {
58+
throw new IllegalArgumentException("Only create table as select statements are supported");
59+
}
60+
61+
final RelRoot queryRelRoot = bestExprRelRootBuilder.apply(converter, sqlCreateTable.query);
62+
return new CreateTableAs(sqlCreateTable.name.names, queryRelRoot);
63+
}
64+
65+
protected CalciteOperation handleCreateView(final SqlCreateView sqlCreateView) {
66+
final RelRoot queryRelRoot = bestExprRelRootBuilder.apply(converter, sqlCreateView.query);
67+
return new CreateView(sqlCreateView.name.names, queryRelRoot);
68+
}
69+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/CalciteOperationVisitor.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package io.substrait.isthmus.operation;
2+
3+
public interface CalciteOperationVisitor<R> {
4+
R visit(CalciteOperation operation);
5+
6+
R visit(RelationalOperation relationalOperation);
7+
8+
R visit(DdlOperation ddlOperation);
9+
10+
R visit(CreateWithInput createWithInput);
11+
12+
R visit(Create create);
13+
14+
R visit(CreateTableAs createTableAs);
15+
16+
R visit(CreateView createView);
17+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/Create.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package io.substrait.isthmus.operation;
2+
3+
import java.util.List;
4+
5+
public abstract class Create extends DdlOperation {
6+
private final List<String> names;
7+
private final RelationType relationType;
8+
9+
public Create(RelationType relationType, List<String> names) {
10+
this.relationType = relationType;
11+
this.names = names;
12+
}
13+
14+
@Override
15+
public List<String> getNames() {
16+
return names;
17+
}
18+
19+
@Override
20+
public RelationType getRelationType() {
21+
return relationType;
22+
}
23+
24+
@Override
25+
public Operation getOperation() {
26+
return Operation.CREATE;
27+
}
28+
29+
@Override
30+
public <R> R accept(CalciteOperationVisitor<R> visitor) {
31+
return visitor.visit(this);
32+
}
33+
}

‎isthmus/src/main/java/io/substrait/isthmus/operation/CreateTableAs.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package io.substrait.isthmus.operation;
2+
3+
import java.util.List;
4+
import org.apache.calcite.rel.RelRoot;
5+
6+
public class CreateTableAs extends CreateWithInput {
7+
public CreateTableAs(List<String> names, RelRoot input) {
8+
super(RelationType.TABLE, names, input);
9+
}
10+
11+
@Override
12+
public <R> R accept(CalciteOperationVisitor<R> visitor) {
13+
return visitor.visit(this);
14+
}
15+
}

0 commit comments

Comments
 (0)