diff --git a/spark/src/main/scala/io/substrait/spark/FileHolder.scala b/spark/src/main/scala/io/substrait/spark/FileHolder.scala index baa4cf5f8..cf774d779 100644 --- a/spark/src/main/scala/io/substrait/spark/FileHolder.scala +++ b/spark/src/main/scala/io/substrait/spark/FileHolder.scala @@ -6,6 +6,7 @@ import io.substrait.relation.{ProtoRelConverter, RelProtoConverter} import io.substrait.relation.Extension.WriteExtensionObject import io.substrait.relation.files.FileOrFiles +/** File target for unpartitioned, unbucketed INSERT writes with append semantics. */ case class FileHolder(fileOrFiles: FileOrFiles) extends WriteExtensionObject { override def toProto(converter: RelProtoConverter): protobuf.Any = { diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala index c3fa30dba..776a65b04 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala @@ -531,12 +531,19 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes } override def visit(write: ExtensionWrite, context: EmptyVisitationContext): LogicalPlan = { - val child = write.getInput.accept(this, context) val mode = write.getOperation match { case WriteOp.INSERT => SaveMode.Append - case WriteOp.UPDATE => SaveMode.Overwrite case op => throw new UnsupportedOperationException(s"Write mode $op not supported") } + // create_mode applies only to CTAS. Older file writes used it for Spark save modes, + // so reject the modes that cannot be represented by an append-only file extension. + write.getCreateMode match { + case CreateMode.UNSPECIFIED | CreateMode.APPEND_IF_EXISTS => + case createMode => + throw new UnsupportedOperationException( + s"Filesystem INSERT does not support create mode $createMode") + } + val child = write.getInput.accept(this, context) val file = write.getDetail match { case FileHolder(f) => f diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala index 90f4a3538..c0facc2fa 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala @@ -606,10 +606,10 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { command match { case InsertIntoHadoopFsRelationCommand( outputPath, - _, - _, - _, - _, + staticPartitions, + ifPartitionNotExists, + partitionColumns, + bucketSpec, fileFormat, options, child, @@ -617,6 +617,16 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { _, _, outputColumnNames) => + if (mode != SaveMode.Append) { + throw new UnsupportedOperationException( + s"Filesystem writes only support SaveMode.Append, found $mode") + } + if (staticPartitions.nonEmpty || ifPartitionNotExists || partitionColumns.nonEmpty) { + throw new UnsupportedOperationException("Partitioned filesystem writes are not supported") + } + if (bucketSpec.nonEmpty) { + throw new UnsupportedOperationException("Bucketed filesystem writes are not supported") + } val file = FileOrFiles .builder() .fileFormat(convertFileFormat(fileFormat, options)) @@ -632,7 +642,7 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { .input(visit(child)) .operation(WriteOp.INSERT) .outputMode(OutputMode.UNSPECIFIED) - .createMode(createMode(mode)) + .createMode(CreateMode.UNSPECIFIED) .tableSchema(outputSchema(child.output, outputColumnNames)) .detail(FileHolder(file)) .build() diff --git a/spark/src/test/scala/io/substrait/spark/FileWriteSuite.scala b/spark/src/test/scala/io/substrait/spark/FileWriteSuite.scala new file mode 100644 index 000000000..e8fc8b788 --- /dev/null +++ b/spark/src/test/scala/io/substrait/spark/FileWriteSuite.scala @@ -0,0 +1,150 @@ +package io.substrait.spark + +import io.substrait.spark.logical.{ToLogicalPlan, ToSubstraitRel} + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.{Row, SaveMode} +import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.catalog.BucketSpec +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.test.SharedSparkSession + +import io.substrait.extension.ExtensionCollector +import io.substrait.relation.{ExtensionWrite, RelProtoConverter} +import io.substrait.relation.AbstractWriteRel.{CreateMode, WriteOp} +import org.apache.hadoop.fs.Path + +class FileWriteSuite extends SparkFunSuite with SharedSparkSession { + + private def withTarget(f: InsertIntoHadoopFsRelationCommand => Unit): Unit = { + withTable("file_write_target") { + spark.sql("CREATE TABLE file_write_target (id INT) USING PARQUET") + spark.sql("INSERT INTO file_write_target VALUES (1), (2)") + val table = spark.sessionState.catalog.getTableMetadata(TableIdentifier("file_write_target")) + val child = spark.sql("SELECT 3 AS id").queryExecution.optimizedPlan + f( + InsertIntoHadoopFsRelationCommand( + outputPath = new Path(table.location), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat(), + options = Map.empty, + query = child, + mode = SaveMode.Append, + catalogTable = Some(table), + fileIndex = None, + outputColumnNames = Seq("id") + )) + } + } + + private def targetRows: Seq[Row] = + spark.sql("SELECT id FROM file_write_target ORDER BY id").collect().toSeq + + private def convertWrite(command: InsertIntoHadoopFsRelationCommand): ExtensionWrite = + new ToSubstraitRel().visit(command).asInstanceOf[ExtensionWrite] + + private def importProto(write: ExtensionWrite): LogicalPlan = { + val collector = new ExtensionCollector + val bytes = new RelProtoConverter(collector).toProto(write).toByteArray + val decoded = new FileHolderHandlingProtoRelConverter(collector) + .from(io.substrait.proto.Rel.parseFrom(bytes)) + new ToLogicalPlan(spark).convert(decoded) + } + + test("append writes preserve existing rows through the file extension protobuf") { + withTarget { + command => + val write = convertWrite(command) + assertResult(CreateMode.UNSPECIFIED)(write.getCreateMode) + val plan = importProto(write) + spark.sessionState.executePlan(plan).executedPlan.execute() + assertResult(Seq(Row(1), Row(2), Row(3)))(targetRows) + } + } + + test("legacy append file extensions remain executable") { + withTarget { + command => + val write = ExtensionWrite + .builder() + .from(convertWrite(command)) + .createMode(CreateMode.APPEND_IF_EXISTS) + .build() + spark.sessionState.executePlan(importProto(write)).executedPlan.execute() + assertResult(Seq(Row(1), Row(2), Row(3)))(targetRows) + } + } + + test("reject filesystem save modes that cannot be represented as INSERT") { + withTarget { + command => + Seq(SaveMode.Overwrite, SaveMode.Ignore, SaveMode.ErrorIfExists).foreach { + mode => + val error = intercept[UnsupportedOperationException] { + convertWrite(command.copy(mode = mode)) + } + assert(error.getMessage.contains(s"SaveMode.Append, found $mode")) + assertResult(Seq(Row(1), Row(2)))(targetRows) + } + } + } + + test("reject partition and bucket metadata that the file extension cannot carry") { + withTarget { + command => + val partitioned = spark.sql("SELECT 3 AS id, 10 AS part").queryExecution.optimizedPlan + val commands = Seq( + command.copy(staticPartitions = Map("part" -> "10")), + command.copy(ifPartitionNotExists = true), + command.copy( + partitionColumns = Seq(partitioned.output.last), + query = partitioned, + outputColumnNames = Seq("id", "part")), + command.copy(bucketSpec = Some(BucketSpec(2, Seq("id"), Seq.empty))) + ) + commands.foreach { + unsupported => + val error = intercept[UnsupportedOperationException] { + convertWrite(unsupported) + } + assert(error.getMessage.contains("filesystem writes are not supported")) + assertResult(Seq(Row(1), Row(2)))(targetRows) + } + } + } + + test("reject legacy file save modes before constructing an executable INSERT") { + withTarget { + command => + Seq(CreateMode.REPLACE_IF_EXISTS, CreateMode.IGNORE_IF_EXISTS, CreateMode.ERROR_IF_EXISTS) + .foreach { + mode => + val write = + ExtensionWrite.builder().from(convertWrite(command)).createMode(mode).build() + val error = intercept[UnsupportedOperationException] { + importProto(write) + } + assert(error.getMessage.contains(s"INSERT does not support create mode $mode")) + assertResult(Seq(Row(1), Row(2)))(targetRows) + } + } + } + + test("reject file UPDATE instead of replacing the entire target") { + withTarget { + command => + val write = + ExtensionWrite.builder().from(convertWrite(command)).operation(WriteOp.UPDATE).build() + val error = intercept[UnsupportedOperationException] { + importProto(write) + } + assert(error.getMessage.contains("Write mode UPDATE not supported")) + assertResult(Seq(Row(1), Row(2)))(targetRows) + } + } +}