diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index 72d319364..2262454a2 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -1071,20 +1071,29 @@ private RelDataType toRowType(NamedStruct schema) { } private RelNode handleCreateTableAs(NamedWrite namedWrite, Context context) { - if (namedWrite.getCreateMode() != AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS - || namedWrite.getOutputMode() != AbstractWriteRel.OutputMode.NO_OUTPUT) { + if (namedWrite.getOutputMode() != AbstractWriteRel.OutputMode.NO_OUTPUT) { throw new UnsupportedOperationException( String.format( - "Can only handle CTAS NamedWrite with (%s, %s), given (%s, %s)", - AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS, - AbstractWriteRel.OutputMode.NO_OUTPUT, - namedWrite.getCreateMode(), - namedWrite.getOutputMode())); + "Can only handle CTAS NamedWrite with output mode %s, given %s", + AbstractWriteRel.OutputMode.NO_OUTPUT, namedWrite.getOutputMode())); + } + switch (namedWrite.getCreateMode()) { + case ERROR_IF_EXISTS: + case IGNORE_IF_EXISTS: + case REPLACE_IF_EXISTS: + break; + default: + throw new UnsupportedOperationException( + "Cannot convert CTAS creation mode to Calcite: " + namedWrite.getCreateMode()); } Rel input = namedWrite.getInput(); RelNode relNode = input.accept(this, context); - return new CreateTable(namedWrite.getNames(), toRowType(namedWrite.getTableSchema()), relNode); + return new CreateTable( + namedWrite.getNames(), + toRowType(namedWrite.getTableSchema()), + relNode, + namedWrite.getCreateMode()); } @Override diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index 479b6fb8f..825419e74 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java @@ -944,7 +944,7 @@ public Rel handleCreateTable(CreateTable createTable) { .input(inputRel) .tableSchema(schema) .operation(AbstractWriteRel.WriteOp.CTAS) - .createMode(AbstractWriteRel.CreateMode.REPLACE_IF_EXISTS) + .createMode(createTable.getCreateMode()) .outputMode(AbstractWriteRel.OutputMode.NO_OUTPUT) .names(createTable.getTableName()) .build(); diff --git a/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/CreateTable.java b/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/CreateTable.java index 6e536417f..3571ec389 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/CreateTable.java +++ b/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/CreateTable.java @@ -1,6 +1,8 @@ package io.substrait.isthmus.calcite.rel; +import io.substrait.relation.AbstractWriteRel.CreateMode; import java.util.List; +import java.util.Objects; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelTraitSet; import org.apache.calcite.rel.RelNode; @@ -13,30 +15,47 @@ public class CreateTable extends SingleRel { private final List tableName; private final RelDataType tableSchema; + private final CreateMode createMode; private CreateTable( RelOptCluster cluster, RelTraitSet traitSet, List tableName, RelDataType tableSchema, - RelNode input) { + RelNode input, + CreateMode createMode) { super(cluster, traitSet, input); this.tableName = tableName; this.tableSchema = DdlSchemas.requireFilledBy(tableSchema, input, "table"); + this.createMode = Objects.requireNonNull(createMode, "createMode"); } /** * CreateTable Constructor, taking the row type of the input as the schema of the table to create. + * Retains the historical replace-if-exists behavior; use the overload with an explicit mode to + * choose another policy. * * @param tableName tablename components * @param input RelNode input */ public CreateTable(List tableName, RelNode input) { - this(input.getCluster(), input.getTraitSet(), tableName, input.getRowType(), input); + this(tableName, input, CreateMode.REPLACE_IF_EXISTS); } /** - * CreateTable Constructor. + * Creates a table with the input's schema and an explicit policy for an existing table. + * + * @param tableName table name components + * @param input the query filling the table + * @param createMode the policy when the target table already exists + */ + public CreateTable(List tableName, RelNode input, CreateMode createMode) { + this(input.getCluster(), input.getTraitSet(), tableName, input.getRowType(), input, createMode); + } + + /** + * CreateTable Constructor. Retains the historical replace-if-exists behavior; use the overload + * with an explicit mode to choose another policy. * * @param tableName tablename components * @param tableSchema the schema of the table to create, which the input fills but need not name @@ -44,7 +63,20 @@ public CreateTable(List tableName, RelNode input) { * @param input RelNode input */ public CreateTable(List tableName, RelDataType tableSchema, RelNode input) { - this(input.getCluster(), input.getTraitSet(), tableName, tableSchema, input); + this(tableName, tableSchema, input, CreateMode.REPLACE_IF_EXISTS); + } + + /** + * Creates a table with a declared schema and an explicit policy for an existing table. + * + * @param tableName table name components + * @param tableSchema the declared schema of the table + * @param input the query filling the table + * @param createMode the policy when the target table already exists + */ + public CreateTable( + List tableName, RelDataType tableSchema, RelNode input, CreateMode createMode) { + this(input.getCluster(), input.getTraitSet(), tableName, tableSchema, input, createMode); } /** @@ -68,7 +100,8 @@ protected RelDataType deriveRowType() { public RelWriter explainTerms(RelWriter pw) { return super.explainTerms(pw) .item("tableName", getTableName()) - .item("tableSchema", getTableSchema().getFullTypeString()); + .item("tableSchema", getTableSchema().getFullTypeString()) + .item("createMode", getCreateMode()); } /** @@ -85,7 +118,8 @@ public RelNode copy(RelTraitSet traitSet, List inputs) { throw new IllegalArgumentException( "CreateTable requires exactly one input, but got " + inputs.size()); } - return new CreateTable(getCluster(), traitSet, tableName, tableSchema, inputs.get(0)); + return new CreateTable( + getCluster(), traitSet, tableName, tableSchema, inputs.get(0), createMode); } /** @@ -107,4 +141,13 @@ public List getTableName() { public RelDataType getTableSchema() { return tableSchema; } + + /** + * Returns the policy to apply when the target table already exists. + * + * @return the creation mode + */ + public CreateMode getCreateMode() { + return createMode; + } } diff --git a/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/DdlSqlToRelConverter.java b/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/DdlSqlToRelConverter.java index 60e6de30d..f18449511 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/DdlSqlToRelConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/calcite/rel/DdlSqlToRelConverter.java @@ -1,5 +1,6 @@ package io.substrait.isthmus.calcite.rel; +import io.substrait.relation.AbstractWriteRel.CreateMode; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -92,18 +93,27 @@ protected RelRoot handleNonDdl(final SqlNode sqlNode) { * * @param sqlCreateTable the CREATE TABLE node * @return a {@link RelRoot} wrapping a synthetic {@code CreateTable} relational node - * @throws IllegalArgumentException if the statement is not a CTAS + * @throws IllegalArgumentException if the statement is not a CTAS or combines OR REPLACE and IF + * NOT EXISTS */ protected RelRoot handleCreateTable(final SqlCreateTable sqlCreateTable) { if (sqlCreateTable.query == null) { throw new IllegalArgumentException("Only create table as select statements are supported"); } + if (sqlCreateTable.getReplace() && sqlCreateTable.ifNotExists) { + throw new IllegalArgumentException( + "CREATE TABLE cannot combine OR REPLACE and IF NOT EXISTS"); + } + final CreateMode createMode = + sqlCreateTable.getReplace() + ? CreateMode.REPLACE_IF_EXISTS + : sqlCreateTable.ifNotExists ? CreateMode.IGNORE_IF_EXISTS : CreateMode.ERROR_IF_EXISTS; final RelNode input = converter.convertQuery(sqlCreateTable.query, true, true).rel; final RelDataType schema = declaredSchema(sqlCreateTable.columnList, input); return RelRoot.of( schema == null - ? new CreateTable(sqlCreateTable.name.names, input) - : new CreateTable(sqlCreateTable.name.names, schema, input), + ? new CreateTable(sqlCreateTable.name.names, input, createMode) + : new CreateTable(sqlCreateTable.name.names, schema, input, createMode), sqlCreateTable.getKind()); } diff --git a/isthmus/src/test/java/io/substrait/isthmus/CtasCreateModeTest.java b/isthmus/src/test/java/io/substrait/isthmus/CtasCreateModeTest.java new file mode 100644 index 000000000..549fd9235 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/CtasCreateModeTest.java @@ -0,0 +1,120 @@ +package io.substrait.isthmus; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.isthmus.calcite.rel.CreateTable; +import io.substrait.isthmus.sql.SubstraitCreateStatementParser; +import io.substrait.isthmus.sql.SubstraitSqlToCalcite; +import io.substrait.plan.Plan; +import io.substrait.plan.PlanProtoConverter; +import io.substrait.plan.ProtoPlanConverter; +import io.substrait.relation.AbstractWriteRel.CreateMode; +import io.substrait.relation.NamedWrite; +import java.util.List; +import org.apache.calcite.prepare.Prepare; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.sql.parser.SqlParseException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.ValueSource; + +class CtasCreateModeTest { + private final ConverterProvider provider = ConverterProvider.DEFAULT; + // An existing target must not turn CREATE or IF NOT EXISTS into a replacement request. + private final Prepare.CatalogReader catalog = + SubstraitCreateStatementParser.processCreateStatementsToCatalog( + provider, "CREATE TABLE dst (v INTEGER)"); + + CtasCreateModeTest() throws SqlParseException {} + + @ParameterizedTest + @CsvSource({ + "CREATE TABLE dst AS SELECT 99 AS v, ERROR_IF_EXISTS", + "CREATE TABLE IF NOT EXISTS dst AS SELECT 99 AS v, IGNORE_IF_EXISTS", + "CREATE OR REPLACE TABLE dst AS SELECT 99 AS v, REPLACE_IF_EXISTS", + "CREATE TABLE dst(v INTEGER) AS SELECT 99 AS v, ERROR_IF_EXISTS", + "CREATE TABLE IF NOT EXISTS dst(v INTEGER) AS SELECT 99 AS v, IGNORE_IF_EXISTS", + "CREATE OR REPLACE TABLE dst(v INTEGER) AS SELECT 99 AS v, REPLACE_IF_EXISTS" + }) + void preservesCreationModeThroughSqlProtoAndCalcite(String sql, CreateMode expected) + throws SqlParseException { + Plan plan = new SqlToSubstrait(provider).convert(sql, catalog); + NamedWrite write = assertInstanceOf(NamedWrite.class, plan.getRoots().get(0).getInput()); + assertEquals(expected, write.getCreateMode()); + + io.substrait.proto.Plan proto = new PlanProtoConverter().toProto(plan); + assertEquals( + expected.toProto(), proto.getRelations(0).getRoot().getInput().getWrite().getCreateMode()); + Plan decoded = new ProtoPlanConverter().from(proto); + CreateTable calcite = + assertInstanceOf( + CreateTable.class, + new SubstraitToCalcite(provider, catalog) + .convert(decoded.getRoots().get(0).getInput())); + assertEquals(expected, calcite.getCreateMode()); + + CreateTable copied = + assertInstanceOf( + CreateTable.class, calcite.copy(calcite.getTraitSet(), List.of(calcite.getInput()))); + NamedWrite roundTripped = + assertInstanceOf(NamedWrite.class, SubstraitRelVisitor.convert(copied, provider)); + assertEquals(expected, roundTripped.getCreateMode()); + } + + @Test + void rejectsConflictingCreationPolicies() { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + new SqlToSubstrait(provider) + .convert( + "CREATE OR REPLACE TABLE IF NOT EXISTS dst AS SELECT 99 AS v", catalog)); + assertTrue(error.getMessage().contains("cannot combine OR REPLACE and IF NOT EXISTS")); + } + + @ParameterizedTest + @ValueSource(strings = {"CREATE TABLE", "CREATE TABLE IF NOT EXISTS", "CREATE OR REPLACE TABLE"}) + void createWithoutAQueryRemainsUnsupported(String prefix) { + assertThrows( + IllegalArgumentException.class, + () -> new SqlToSubstrait(provider).convert(prefix + " dst(v INTEGER)", catalog)); + } + + @Test + void existingConstructorsRetainTheirCreationPolicy() throws SqlParseException { + RelNode input = SubstraitSqlToCalcite.convertQuery("SELECT 99 AS v", catalog, provider).rel; + assertEquals( + CreateMode.REPLACE_IF_EXISTS, new CreateTable(List.of("DST"), input).getCreateMode()); + assertEquals( + CreateMode.REPLACE_IF_EXISTS, + new CreateTable(List.of("DST"), input.getRowType(), input).getCreateMode()); + } + + @Test + void differentCreationPoliciesHaveDifferentPlannerDigests() throws SqlParseException { + RelNode input = SubstraitSqlToCalcite.convertQuery("SELECT 99 AS v", catalog, provider).rel; + CreateTable plain = new CreateTable(List.of("DST"), input, CreateMode.ERROR_IF_EXISTS); + CreateTable replace = new CreateTable(List.of("DST"), input, CreateMode.REPLACE_IF_EXISTS); + assertNotEquals(plain.getDigest(), replace.getDigest()); + } + + @ParameterizedTest + @EnumSource( + value = CreateMode.class, + names = {"UNSPECIFIED", "APPEND_IF_EXISTS"}) + void unsupportedCreationModesAreRefused(CreateMode mode) throws SqlParseException { + Plan plan = new SqlToSubstrait(provider).convert("CREATE TABLE dst AS SELECT 99 AS v", catalog); + NamedWrite write = assertInstanceOf(NamedWrite.class, plan.getRoots().get(0).getInput()); + NamedWrite unsupported = NamedWrite.builder().from(write).createMode(mode).build(); + assertThrows( + UnsupportedOperationException.class, + () -> new SubstraitToCalcite(provider, catalog).convert(unsupported)); + } +}