diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 69f0fa89612..debfccf311e 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -460,6 +460,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 0b525419943..742fd100d09 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -164,6 +164,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase diff --git a/.github/workflows/spark_sql_writer_tests.yml b/.github/workflows/spark_sql_writer_tests.yml index b26ff88e864..6b5fc7f6d65 100644 --- a/.github/workflows/spark_sql_writer_tests.yml +++ b/.github/workflows/spark_sql_writer_tests.yml @@ -115,6 +115,13 @@ jobs: spark-short-version: ${{ inputs.spark-version }} skip-native-build: true + - name: Pre-compile Spark SQL test classes + run: | + cd apache-spark + rm -rf /root/.m2/repository/org/apache/parquet + NOLINT_ON_COMPILE=true build/sbt -Dsbt.log.noformat=true -mem 3072 \ + 'sql/Test/compile' + - name: Run Parquet writer tests run: | cd apache-spark diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 4c65b2ee6c5..40732ced79f 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -29,8 +29,10 @@ executed. Supported parent joins and aggregates remain eligible for native execution. Global aggregates still return one row (`COUNT = 0`, `SUM = NULL`), and grouped aggregates return no rows. Independent operator restrictions and aggregate buffer compatibility checks still apply. -Parquet writes whose input plans contain an empty relation use Spark's writer to preserve -readable empty output files and their schema metadata. + +A native Parquet write over a native empty relation stays native. Like Spark's writer, it runs one +task for the empty input, so the output still gets a schema-only Parquet file that readers can +infer the schema from. ## In-Memory Cache diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 0155ebb8dae..acefd4f0890 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -64,7 +64,7 @@ omitted from the tables below and may be reconsidered based on demand: | `FileSourceScanExec` | ✅ | Parquet only. Some types and configurations fall back. See [Parquet Scan Compatibility](compatibility/scans.md). | | `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet Scan Compatibility](compatibility/scans.md) and the [Iceberg Guide](iceberg.md). | | `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). | -| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support and writer fallback. | +| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support, including native Parquet writes. | | `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. | | `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index c96cb17020e..cc2da616e54 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -234,18 +234,9 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Where this task writes. When `work_dir` is set (the Spark 3.x `CometNativeWriteExec` - /// path) this is the write's output directory and is unused; the file name is derived from - /// `work_dir`. Otherwise (Spark 4.0+, `CometWriteFilesExec`) it is the exact path of the file - /// to write, chosen by the JVM commit protocol and used verbatim - this operator then never - /// derives file names of its own. + /// The exact path of the file this task writes, chosen by Spark's commit protocol on the JVM + /// side and used verbatim - this operator never derives file names of its own. output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only. - work_dir: Option, - /// Job ID for tracking this write operation - job_id: Option, - /// Task attempt ID for this specific task - task_attempt_id: Option, /// Compression codec compression: ParquetCompression, /// Partition ID (from Spark TaskContext) @@ -268,9 +259,6 @@ impl ParquetWriterExec { pub fn try_new( input: Arc, output_path: String, - work_dir: Option, - job_id: Option, - task_attempt_id: Option, compression: ParquetCompression, partition_id: i32, column_names: Vec, @@ -290,9 +278,6 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, output_path, - work_dir, - job_id, - task_attempt_id, compression, partition_id, column_names, @@ -478,9 +463,6 @@ impl ExecutionPlan for ParquetWriterExec { 1 => Ok(Arc::new(ParquetWriterExec::try_new( Arc::clone(&children[0]), self.output_path.clone(), - self.work_dir.clone(), - self.job_id.clone(), - self.task_attempt_id, self.compression.clone(), self.partition_id, self.column_names.clone(), @@ -510,8 +492,6 @@ impl ExecutionPlan for ParquetWriterExec { .register(context.memory_pool()); let input = self.input.execute(partition, context)?; let input_schema = self.input.schema(); - let work_dir = self.work_dir.clone(); - let task_attempt_id = self.task_attempt_id; let compression = self.compression.to_parquet()?; let column_names = self.column_names.clone(); @@ -530,19 +510,8 @@ impl ExecutionPlan for ParquetWriterExec { Arc::new(Schema::new(fields)) }); - let part_file = match &work_dir { - // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. - None => self.output_path.clone(), - // Spark 3.x hands over a working directory instead and expects the writer to name the - // file; that branch goes away with Spark 3.x support. - Some(work_dir) => match task_attempt_id { - Some(attempt_id) => format!( - "{}/part-{:05}-{:05}.parquet", - work_dir, self.partition_id, attempt_id - ), - None => format!("{}/part-{:05}.parquet", work_dir, self.partition_id), - }, - }; + // The JVM commit protocol has already chosen the exact file to write. + let part_file = self.output_path.clone(); // Configure writer properties let props = WriterProperties::builder() @@ -675,11 +644,10 @@ mod tests { ); } - /// Spark 4.0+ hands over the exact file to write rather than a working directory. The writer - /// must use that path verbatim - Spark's commit protocol owns naming and staging, and + /// The JVM hands over the exact file to write. The writer must use that path verbatim - Spark's commit protocol owns naming and staging, and /// committers that track individual files depend on the name it chose. #[tokio::test] - async fn test_parquet_writer_uses_output_path_verbatim_without_work_dir() -> Result<()> { + async fn test_parquet_writer_uses_output_path_verbatim() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); let batch = RecordBatch::try_new( Arc::clone(&schema), @@ -696,9 +664,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, output_path, - None, // work_dir: Spark 4.0+ path - None, - None, ParquetCompression::None, // A non-zero partition id must not leak into the file name. 3, @@ -763,13 +728,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["required_id".to_string(), "values".to_string()], @@ -832,13 +794,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["values".to_string()], @@ -903,9 +862,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, format!("file://{}", file.display()), - None, - None, - None, ParquetCompression::None, 0, column_names, @@ -1160,16 +1116,14 @@ mod tests { let memory_exec = Arc::new(DataSourceExec::new(Arc::new(memory_source_config))); // Create ParquetWriterExec with DataSourceExec as input - let output_path = "unused".to_string(); - let work_dir = "hdfs://namenode:9000/user/test_parquet_writer_exec".to_string(); + let output_path = + "hdfs://namenode:9000/user/test_parquet_writer_exec/part-00000-00123.parquet" + .to_string(); let column_names = vec!["id".to_string(), "name".to_string()]; let parquet_writer = ParquetWriterExec::try_new( memory_exec, output_path, - Some(work_dir), - None, // job_id - Some(123), // task_attempt_id ParquetCompression::None, 0, // partition_id column_names, diff --git a/native/core/src/execution/planner/write.rs b/native/core/src/execution/planner/write.rs index 1a3f9674cd1..5665939afb2 100644 --- a/native/core/src/execution/planner/write.rs +++ b/native/core/src/execution/planner/write.rs @@ -107,9 +107,6 @@ impl OperatorBuilder for ParquetWriterBuilder { let parquet_writer = Arc::new(ParquetWriterExec::try_new( Arc::clone(&child.native_plan), writer.output_path.clone(), - writer.work_dir.clone(), - writer.job_id.clone(), - writer.task_attempt_id, codec, planner.partition(), writer.column_names.clone(), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 527461f47a5..97c65f0da36 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -785,7 +785,7 @@ message IcebergWriteCommon { } // Single Iceberg write operator. Per-task fields are populated by the JVM exec wrapper at task -// launch (matching how `ParquetWriter.task_attempt_id` is filled in `CometNativeWriteExec`). +// launch (matching how `ParquetWriter.output_path` is filled in by the Parquet write execs). message IcebergWrite { IcebergWriteCommon common = 1; @@ -920,27 +920,19 @@ message ShuffleWriter { } message ParquetWriter { - // Where this task writes. Two shapes, selected by whether `work_dir` is set: - // - // - Spark 4.0+ (CometWriteFilesExec): `work_dir` is unset and `output_path` is the - // fully-qualified path of the Parquet file to write, set per task from - // FileCommitProtocol.newTaskTempFile. Naming and staging are owned by Spark's commit - // protocol so that task-attempt isolation, speculative execution, and committers that track - // individual files (S3A magic, streaming manifest) all behave as they do for Spark's own - // writer. The native writer uses this path verbatim. - // - Spark 3.x (CometNativeWriteExec): `work_dir` is set and the native writer derives the file - // name from it, the partition id and the task attempt id. `output_path` is the write's - // output directory and is unused natively. Goes away with Spark 3.x support. + // The fully-qualified path of the Parquet file this task writes, set per task from + // FileCommitProtocol.newTaskTempFile by CometWriteFilesExec (Spark 4.0+) or + // CometNativeWriteExec (Spark 3.x). Naming and staging are owned by Spark's commit protocol so + // that task-attempt isolation, speculative execution, and committers that track individual + // files (S3A magic, streaming manifest) all behave as they do for Spark's own writer. The native + // writer uses this path verbatim. string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; - // Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only; see - // output_path above. - optional string work_dir = 5; - // Job ID for tracking this write operation - optional string job_id = 6; - // Task attempt ID for this specific task - optional int32 task_attempt_id = 7; + // Formerly work_dir, job_id and task_attempt_id, which let the native writer name its own file + // under a working directory. + reserved 5, 6, 7; + reserved "work_dir", "job_id", "task_attempt_id"; // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object // stores. diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index a56dea862c4..545cbc774ae 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -19,17 +19,21 @@ package org.apache.comet.serde.operator +import java.util.UUID + import scala.jdk.CollectionConverters._ -import org.apache.spark.SparkException -import org.apache.spark.sql.comet.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.hadoop.mapreduce.Job +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.spark.internal.io.FileCommitProtocol +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec} import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.util.SerializableConfiguration import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -55,25 +59,18 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec case cmd: InsertIntoHadoopFsRelationCommand => cmd.fileFormat match { case _: ParquetFileFormat => - // AQE can replace the write input with a zero-partition empty relation. Keep - // Spark's writer, which creates an empty task to preserve the output file schema. - // The native writer only maps existing partitions; see #5303. This guard is - // conservative: an empty relation below an exchange can have nonzero partitions - // at the write input. Revisit the guard when native empty-file handling is fixed. - if (hasEmptyRelationInput(op.child)) { - return Unsupported(Some( - "Parquet writes with empty-relation inputs require Spark's empty-file handling")) - } - if (!cmd.outputPath.toString.startsWith("file:") && !cmd.outputPath.toString .startsWith("hdfs:")) { return Unsupported(Some("Supported output filesystems: local, HDFS")) } + // Spark 3.x HadoopMapReduceCommitProtocol.getFilename hardcodes "part"; + // mapreduce.output.basename is only honored on Spark 4.0+. Custom committer paths + // are still checked by checkNativeWriteDestination before native execution. NativeWriteUtils - // This writer names its own files `part--.parquet`, so the - // prefix is fixed rather than read from `mapreduce.output.basename`. - .escapedHdfsDestination(cmd.outputPath.toString, "part") + .escapedHdfsDestination( + cmd.outputPath.toString, + NativeWriteUtils.DEFAULT_BASE_OUTPUT_NAME) .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -98,14 +95,6 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec } } - private def hasEmptyRelationInput(plan: SparkPlan): Boolean = plan match { - case _: CometEmptyRelationExec => true - case wrapper: CometScanWrapper => hasEmptyRelationInput(wrapper.originalPlan) - case stage: QueryStageExec => hasEmptyRelationInput(stage.plan) - case reused: ReusedExchangeExec => hasEmptyRelationInput(reused.child) - case _ => plan.children.exists(hasEmptyRelationInput) - } - override def convert( op: DataWritingCommandExec, builder: Operator.Builder, @@ -123,6 +112,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec val outputPath = cmd.outputPath.toString + // Planning-time value only. CometNativeWriteExec replaces it per task with the codec that + // names the file. val plannedCodec = NativeWriteUtils.parseCompressionCodec(cmd.options) val codec = NativeWriteUtils.protoCompressionCodec(plannedCodec) match { case Some(codec) => codec @@ -140,8 +131,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec cmd.query.schema.fields.toIndexedSeq, Some( op.session.sessionState.conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))).asJava) - // Note: work_dir, job_id, and task_attempt_id will be set at execution time - // in CometNativeWriteExec, as they depend on the Spark task context + // CometNativeWriteExec replaces output_path with the committer's exact task filename + // at execution time. // Collect S3/cloud storage configurations val session = op.session @@ -189,29 +180,30 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec other } - // Create FileCommitProtocol for atomic writes - val jobId = java.util.UUID.randomUUID().toString - val committer = - try { - // Use Spark's SQLHadoopMapReduceCommitProtocol - val committerClass = - classOf[org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol] - val constructor = - committerClass.getConstructor(classOf[String], classOf[String], classOf[Boolean]) - Some( - constructor - .newInstance( - jobId, - outputPath, - java.lang.Boolean.FALSE // dynamicPartitionOverwrite = false for now - ) - .asInstanceOf[org.apache.spark.internal.io.FileCommitProtocol]) - } catch { - case e: Exception => - throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}") - } - - CometNativeWriteExec(nativeOp, op, childPlan, outputPath, cmd.mode, committer, jobId) + val session = op.session + val job = Job.getInstance(session.sessionState.newHadoopConfWithOptions(cmd.options)) + job.setOutputKeyClass(classOf[Void]) + job.setOutputValueClass(classOf[InternalRow]) + FileOutputFormat.setOutputPath(job, cmd.outputPath) + val outputWriterFactory = + cmd.fileFormat.prepareWrite(session, job, CaseInsensitiveMap(cmd.options), cmd.query.schema) + + val committer = FileCommitProtocol.instantiate( + session.sessionState.conf.fileCommitProtocolClass, + UUID.randomUUID().toString, + outputPath, + dynamicPartitionOverwrite = false) + job.getConfiguration.set("spark.sql.sources.writeJobUUID", UUID.randomUUID().toString) + + CometNativeWriteExec( + nativeOp, + op, + childPlan, + outputPath, + cmd.mode, + committer, + new SerializableConfiguration(job.getConfiguration), + outputWriterFactory) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 5246099fdcb..a0ebdea520f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -141,12 +141,13 @@ object NativeWriteUtils { * byte for byte (see [[hdfsPathDivergence]] for why it may not): * * - the destination directory, and - * - `fileNamePrefix`, the basename every file name is built from. On Spark 4.0+ that is - * `mapreduce.output.basename`, which `HadoopMapReduceCommitProtocol.getFilename` - * interpolates into `--`; on 3.x Comet names the files itself and - * the basename is always the literal `part`. A basename holding `?` or `#` is the dangerous - * one: the native URL parser truncates there, so *every* task writes a file with the same - * truncated name and they overwrite each other during commit. + * - `fileNamePrefix`, the basename every file name is built from. Spark 4.0+ + * `HadoopMapReduceCommitProtocol.getFilename` reads `mapreduce.output.basename` and + * interpolates it into `--`; Spark 3.x hardcodes `part` instead. + * Both native writers take the exact file name from the commit protocol. A basename holding + * `?` or `#` is dangerous on 4.0+: the native URL parser truncates there, so every task + * writes a file with the same truncated name and they overwrite each other during commit. + * Custom committer paths are checked by [[checkNativeWriteDestination]]. * * The basename is checked by running [[hdfsPathDivergence]] over the path it produces rather * than over the name alone. Everything else `getFilename` interpolates -- the split number, the diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index bb2c560ba64..200c7831321 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -19,13 +19,16 @@ package org.apache.spark.sql.comet -import scala.jdk.CollectionConverters._ +import java.util.Date +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext, TaskAttemptID, TaskID, TaskType} import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.parquet.hadoop.codec.CodecConfig import org.apache.spark.TaskContext -import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec} +import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage import org.apache.spark.rdd.RDD import org.apache.spark.sql.SaveMode import org.apache.spark.sql.catalyst.InternalRow @@ -35,42 +38,23 @@ import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.command.DataWritingCommandExec -import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.{OutputWriterFactory, WriteFilesExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils +import org.apache.spark.util.{SerializableConfiguration, Utils} -import com.google.protobuf.CodedOutputStream - -import org.apache.comet.CometExecIterator import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.NativeWriteUtils /** - * Comet physical operator for native Parquet write operations with FileCommitProtocol support. - * - * This operator writes data to Parquet files using the native Comet engine. It integrates with - * Spark's FileCommitProtocol to provide atomic writes with proper staging and commit semantics. + * Native Parquet writes on Spark 3.x, where Comet replaces the entire write command. * - * The implementation includes support for Spark's file commit protocol through work_dir, job_id, - * and task_attempt_id parameters that can be set in the operator. When work_dir is set, files are - * written to a temporary location that can be atomically committed later. + * Both execution entry points run Spark's job and task commit lifecycle. Native code writes the + * exact filename supplied by the configured commit protocol, using the prepared Hadoop job. Spark + * 4.0+ uses CometWriteFilesExec and leaves the surrounding command with Spark instead. * - * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) * @param originalPlan * The JVM data-writing command restored when Comet reverts a transition-heavy stage - * @param child - * The child operator providing the data to write - * @param outputPath - * The path where the Parquet file will be written - * @param mode - * The Spark SaveMode governing target-exists behavior (Append / Overwrite / ErrorIfExists / - * Ignore). Comet takes over Spark's DataWritingCommandExec so we must apply these semantics - * here - a direct port of Spark's InsertIntoHadoopFsRelationCommand.run() doInsertion logic. - * @param committer - * FileCommitProtocol for atomic writes. If None, files are written directly. - * @param jobTrackerID - * Unique identifier for this write job */ case class CometNativeWriteExec( nativeOp: Operator, @@ -78,8 +62,10 @@ case class CometNativeWriteExec( child: SparkPlan, outputPath: String, mode: SaveMode, - committer: Option[FileCommitProtocol] = None, - jobTrackerID: String = Utils.createTempDir().getName) + committer: FileCommitProtocol, + serializableHadoopConf: SerializableConfiguration, + outputWriterFactory: OutputWriterFactory, + jobTrackerID: String = SparkHadoopWriterUtils.createJobTrackerID(new Date())) extends CometNativeExec with UnaryExecNode { @@ -97,19 +83,8 @@ case class CometNativeWriteExec( s"${getClass.getSimpleName} expected one reverted input but received ${newChildren.size}") } - // Accumulator to collect TaskCommitMessages from all tasks - // Must be eagerly initialized on driver, not lazy - @transient private val taskCommitMessagesAccum = - sparkContext.collectionAccumulator[FileCommitProtocol.TaskCommitMessage]("taskCommitMessages") - - override def serializedPlanOpt: SerializedPlan = { - val size = nativeOp.getSerializedSize - val bytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(bytes) - nativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - SerializedPlan(Some(bytes)) - } + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) @@ -122,219 +97,132 @@ case class CometNativeWriteExec( "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) override def doExecute(): RDD[InternalRow] = { - // Setup job if committer is present - committer.foreach { c => - val jobContext = createJobContext() - c.setupJob(jobContext) - } - - // Execute the native write with commit protocol - val resultRDD = doExecuteColumnar() - - // Force execution by consuming all batches - resultRDD - .mapPartitions { iter => - iter.foreach(_.close()) - Iterator.empty - } - .count() - - // Extract write statistics from metrics - val filesWritten = metrics("files_written").value - val bytesWritten = metrics("bytes_written").value - val rowsWritten = metrics("rows_written").value - - // Collect TaskCommitMessages from accumulator - val commitMessages = taskCommitMessagesAccum.value.asScala.toSeq - - // Commit job with collected TaskCommitMessages - committer.foreach { c => - val jobContext = createJobContext() - try { - c.commitJob(jobContext, commitMessages) - logInfo( - s"Successfully committed write job to $outputPath: " + - s"$filesWritten files, $bytesWritten bytes, $rowsWritten rows") - } catch { - case e: Exception => - logError("Failed to commit job, aborting", e) - c.abortJob(jobContext) - throw e - } - } - - // Return empty RDD as write operations don't return data + executeWriteAndCommit() sparkContext.emptyRDD[InternalRow] } override def doExecuteColumnar(): RDD[ColumnarBatch] = { - // Comet replaces DataWritingCommandExec entirely, so Spark's - // InsertIntoHadoopFsRelationCommand.run() never runs. That method is where Spark handles - // SaveMode semantics (path-exists check, delete-before-Overwrite, Ignore short-circuit) - - // port the non-partitioned, non-catalog branch of that logic here. See Spark 3.5's - // InsertIntoHadoopFsRelationCommand.run doInsertion match. This runs on the driver before - // any executor tasks fire, mirroring where Spark does the delete. + executeWriteAndCommit() + sparkContext.emptyRDD[ColumnarBatch] + } + + private def executeWriteAndCommit(): Unit = { if (!prepareOutputPathForMode()) { logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") - return sparkContext.emptyRDD[ColumnarBatch] + return } - // Get the input data from the child operator - val childRDD = if (child.supportsColumnar) { + val job = Job.getInstance(new Configuration(serializableHadoopConf.value)) + // Like FileFormatWriter, only abort after setupJob has succeeded. + committer.setupJob(job) + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Include configuration changes made by setupJob in the task contexts. + val commitMessages = runNativeWriteJob(new SerializableConfiguration(job.getConfiguration)) + committer.commitJob(job, commitMessages.toSeq) + logInfo( + s"Successfully committed native write job to $outputPath: " + + s"${metrics("files_written").value} files, " + + s"${metrics("bytes_written").value} bytes, ${metrics("rows_written").value} rows") + })(catchBlock = committer.abortJob(job)) + } + + private def runNativeWriteJob( + hadoopConf: SerializableConfiguration): Array[TaskCommitMessage] = { + val childRDD: RDD[ColumnarBatch] = if (child.supportsColumnar) { child.executeColumnar() } else { - // If child doesn't support columnar, convert to columnar child.execute().mapPartitionsInternal { _ => - // TODO this could delegate to CometRowToColumnar, but maybe Comet - // does not need to support this case? throw new UnsupportedOperationException( "Row-based child operators not yet supported for native write") } } - // Capture metadata before the transformation - val numPartitions = childRDD.getNumPartitions - val numOutputCols = child.output.length + // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, + // so the output directory would carry no schema for readers. Swap in a single empty + // partition, as Spark's FileFormatWriter and CometWriteFilesExec do. + val writeRDD = if (childRDD.getNumPartitions == 0) { + sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) + } else { + childRDD + } + + // The task closure below must not touch `this`, which holds `nativeOp` plus the whole + // converted child subtree; capture what it needs by value instead, as CometWriteFilesExec + // does. + val numPartitions = writeRDD.getNumPartitions + val childSchema = CometUtils.fromAttributes(child.output) + val capturedNativeOp = nativeOp val capturedCommitter = committer val capturedJobTrackerID = jobTrackerID - val capturedNativeOp = nativeOp - val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks - - // Execute native write operation with task-level commit protocol - childRDD.mapPartitionsInternal { iter => - val partitionId = org.apache.spark.TaskContext.getPartitionId() - val taskAttemptId = org.apache.spark.TaskContext.get().taskAttemptId() - - // Setup task-level commit protocol if provided - val (workDir, taskContext, _) = capturedCommitter - .map { committer => - val taskContext = - createTaskContext(capturedJobTrackerID, partitionId, taskAttemptId.toInt) - - // Setup task - this creates the temporary working directory - committer.setupTask(taskContext) - - // Get the work directory for temp files - // Spark 4.1 made the (taskContext, dir, ext: String) overload throw by default; - // the FileNameSpec overload is the supported one and exists in 3.4+. - val workPath = committer.newTaskTempFile(taskContext, None, FileNameSpec("", "")) - val workDir = new Path(workPath).getParent.toString - - (Some(workDir), Some((committer, taskContext)), null) - } - .getOrElse((None, None, null)) - - // Modify the native operator to include task-specific parameters - val modifiedNativeOp = if (workDir.isDefined) { - val parquetWriter = capturedNativeOp.getParquetWriter.toBuilder - .setWorkDir(workDir.get) - .setJobId(capturedJobTrackerID) - .setTaskAttemptId(taskAttemptId.toInt) - .build() - - capturedNativeOp.toBuilder.setParquetWriter(parquetWriter).build() - } else { - capturedNativeOp - } - - val nativeMetrics = CometMetricNode.fromCometPlan(this) - // Register before CometExecIterator so completion listeners run after iterator close - // (Spark runs task completion callbacks in reverse registration order). - Option(TaskContext.get()).foreach(nativeMetrics.reportNativeWriteOutputMetrics) - - val size = modifiedNativeOp.getSerializedSize - val planBytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(planBytes) - modifiedNativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - - val execIterator = new CometExecIterator( - CometExec.newIterId, - CometArrowStream.inputObjects( - iter, - CometUtils.fromAttributes(child.output), - "CometNativeWriteExec"), - numOutputCols, - planBytes, - nativeMetrics, - numPartitions, - partitionId, - None, - Seq.empty) - - // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage - new Iterator[ColumnarBatch] { - private var completed = false - private var thrownException: Option[Throwable] = None - - override def hasNext: Boolean = { - val result = - try { - execIterator.hasNext - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - - if (!result && !completed) { - handleTaskEnd() - } - - result - } - - override def next(): ColumnarBatch = { - try { - execIterator.next() - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - } - - private def handleTaskEnd(): Unit = { - if (!completed) { - completed = true - - // Handle commit or abort based on whether an exception was thrown - taskContext.foreach { case (committer, ctx) => - try { - if (thrownException.isEmpty) { - // Commit the task and add message to accumulator - val message = committer.commitTask(ctx) - capturedAccumulator.add(message) - logDebug(s"Task ${ctx.getTaskAttemptID} committed successfully") - } else { - // Abort the task - committer.abortTask(ctx) - val exMsg = thrownException.get.getMessage - logWarning(s"Task ${ctx.getTaskAttemptID} aborted due to exception: $exMsg") - } - } catch { - case e: Exception => - // Log the commit/abort exception but don't mask the original exception - logError(s"Error during task commit/abort: ${e.getMessage}", e) - if (thrownException.isEmpty) { - // If no original exception, propagate the commit/abort exception - throw e - } + val writerFactory = outputWriterFactory + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val commitMessages = new Array[TaskCommitMessage](numPartitions) + + sparkContext.runJob( + writeRDD, + (context: TaskContext, batches: Iterator[ColumnarBatch]) => { + val taskContext = + CometNativeWriteExec.createTaskContext(hadoopConf.value, capturedJobTrackerID, context) + capturedCommitter.setupTask(taskContext) + // Guard filename allocation, native iterator construction, execution, cleanup and commit. + // Spark's helper preserves the original error if abortTask also fails. + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Mirrors FileFormatWriter's EmptyDirectoryDataWriter case: an empty input still writes + // one file from partition 0 so that the output carries the schema, but every other + // empty partition produces no file at all. + if (context.partitionId() == 0 || batches.hasNext) { + val extension = writerFactory.getFileExtension(taskContext) + val filePath = capturedCommitter.newTaskTempFile( + taskContext, + None, + FileNameSpec("", "-c000" + extension)) + NativeWriteUtils.checkNativeWriteDestination(filePath) + // `extension` names the codec in the job configuration. The writer uses that codec + // too, so the file name and the footer agree. + val codec = CodecConfig.from(taskContext).getCodec + val protoCodec = NativeWriteUtils + .protoCompressionCodec(codec.name()) + .getOrElse(throw new UnsupportedOperationException( + s"Comet's native Parquet writer cannot write $codec")) + val writer = capturedNativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .setCompression(protoCodec) + .build() + val taskOp = capturedNativeOp.toBuilder.setParquetWriter(writer).build() + + // Register before the iterator so this listener runs after its cleanup. + nativeMetrics.reportNativeWriteOutputMetrics(context) + val execIterator = CometExec.getCometIterator( + CometArrowStream.inputObjects(batches, childSchema, "CometNativeWriteExec"), + childSchema.length, + taskOp, + nativeMetrics, + numPartitions, + context.partitionId(), + None, + Seq.empty) + + // Close before committing. A failed write must remain the primary error even if + // native teardown (including the final metrics update) also throws. + Utils.tryWithSafeFinally { + while (execIterator.hasNext) { + execIterator.next().close() } + } { + execIterator.close() } } - } - } - } - } - - /** Create a JobContext for the write job */ - private def createJobContext(): Job = { - val job = Job.getInstance() - job.setJobID(new org.apache.hadoop.mapreduce.JobID(jobTrackerID, 0)) - job + // `hasNext` already ran an empty partition's child to completion, so there is nothing + // to drain or release when no file was written. + capturedCommitter.commitTask(taskContext) + })(catchBlock = capturedCommitter.abortTask(taskContext)) + }, + writeRDD.partitions.indices, + (index, message: TaskCommitMessage) => { + committer.onTaskCommit(message) + commitMessages(index) = message + }) + commitMessages } /** @@ -347,7 +235,7 @@ case class CometNativeWriteExec( */ private def prepareOutputPathForMode(): Boolean = { val path = new Path(outputPath) - val hadoopConf = sparkContext.hadoopConfiguration + val hadoopConf = serializableHadoopConf.value val fs = path.getFileSystem(hadoopConf) val qualifiedOutputPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) @@ -361,11 +249,7 @@ case class CometNativeWriteExec( true case SaveMode.Overwrite => if (fs.exists(qualifiedOutputPath)) { - val deleted = committer match { - case Some(c) => c.deleteWithJob(fs, qualifiedOutputPath, true) - case None => fs.delete(qualifiedOutputPath, true) - } - if (!deleted) { + if (!committer.deleteWithJob(fs, qualifiedOutputPath, true)) { throw QueryExecutionErrors.cannotClearOutputDirectoryError(qualifiedOutputPath) } } @@ -375,15 +259,27 @@ case class CometNativeWriteExec( } } - /** Create a TaskAttemptContext for a specific task */ +} + +object CometNativeWriteExec { + + /** + * Match FileFormatWriter's Hadoop task identifiers and configuration. Kept off the exec node so + * that the task closure does not capture it. + */ private def createTaskContext( - jobId: String, - partitionId: Int, - attemptNumber: Int): TaskAttemptContext = { - val job = Job.getInstance() - val taskAttemptID = new TaskAttemptID( - new TaskID(new org.apache.hadoop.mapreduce.JobID(jobId, 0), TaskType.REDUCE, partitionId), - attemptNumber) - new TaskAttemptContextImpl(job.getConfiguration, taskAttemptID) + conf: Configuration, + jobTrackerID: String, + context: TaskContext): TaskAttemptContext = { + val hadoopConf = new Configuration(conf) + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, context.stageId()) + val taskId = new TaskID(jobId, TaskType.MAP, context.partitionId()) + val attemptId = new TaskAttemptID(taskId, context.taskAttemptId().toInt & Integer.MAX_VALUE) + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskId.toString) + hadoopConf.set("mapreduce.task.attempt.id", attemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + new TaskAttemptContextImpl(hadoopConf, attemptId) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala index 305593e609d..7721cb7c2a5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -129,10 +129,10 @@ case class CometWriteFilesExec( // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, // so the output directory would carry no schema for readers. Spark's own WriteFilesExec swaps - // in a dummy single-partition RDD for exactly this case. AQE reaches it by collapsing a - // completed empty shuffle into a CometEmptyRelationExec, which is why - // CometDataWritingCommand declines empty-relation inputs on the Spark 3.x path (#5303): its - // writer has nowhere to put this swap. See CometEmptyRelationParquetWriterSuite. + // in a dummy single-partition RDD for exactly this case, and CometNativeWriteExec does the + // same on the Spark 3.x path (#5303). AQE reaches it by collapsing a completed empty shuffle + // into a CometEmptyRelationExec, so a write over a native empty relation stays native and + // still leaves a schema-only file. See CometEmptyRelationParquetWriterSuite. val writeRDD = if (childRDD.getNumPartitions == 0) { sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) } else { diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index 5b01a603534..466f6cad9e6 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -613,13 +613,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { checkAnswer(spark.read.parquet(outputPath), df.collect()) assertParquetCodec(outputPath, CompressionCodecName.GZIP) - if (isSpark40Plus) { - // Spark names the file; Comet fills it. The extension is the only externally visible - // statement of the codec, so it has to agree with the footer. (On 3.x the native writer - // invents a name with no codec suffix, so there is nothing to compare.) - listPartFileNames(outputPath).foreach { name => - assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") - } + // Spark names the file; Comet fills it. The extension is the only externally visible + // statement of the codec, so it has to agree with the footer. + listPartFileNames(outputPath).foreach { name => + assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") } } } @@ -1237,16 +1234,13 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } // --------------------------------------------------------------------------------------------- - // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in - // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the - // `WriteFilesExecBase` trait. See CometWriteFilesExec. + // Commit-protocol checks run on both writers. Tests requiring the surrounding Spark write + // command are gated to Spark 4.0+, where Comet replaces only WriteFilesExec. // --------------------------------------------------------------------------------------------- test("write creates a _SUCCESS marker") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/2985 - the marker comes from - // HadoopMapReduceCommitProtocol.commitJob, which only runs because Comet leaves - // InsertIntoHadoopFsRelationCommand in the plan. + // HadoopMapReduceCommitProtocol.commitJob on both native writer paths. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath withTempPath { srcDir => @@ -1266,7 +1260,6 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("written file names follow Spark's naming convention") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // The file name comes from FileCommitProtocol.newTaskTempFile and must be used verbatim: // part---c..parquet. Committers that track individual files // and tools that parse these names depend on it. @@ -1295,12 +1288,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } - test("a custom output basename is honored on local storage") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // `escapedHdfsDestination` gates the basename on HDFS only, where the native URL parser would - // rename the file out from under the committer. Local writes hand the path to the native - // writer verbatim, so this is the control that keeps that guard from being widened: `part%foo` - // is declined on HDFS and has to keep working here. + test("local output basenames match Spark's version-specific commit protocol") { + // Spark 4.0+ honors the basename, including literal '%' on local storage where the path + // reaches the native writer verbatim. Spark 3.x ignores the option and always uses "part". Seq("out", "part%foo").foreach { basename => withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath @@ -1316,9 +1306,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } val written = new File(outputPath).listFiles().map(_.getName).filter(_.endsWith(".parquet")) + val expectedBasename = if (isSpark40Plus) basename else "part" assert( - written.nonEmpty && written.forall(_.startsWith(s"$basename-")), - s"expected every data file to be named '$basename-...', found: " + + written.nonEmpty && written.forall(_.startsWith(s"$expectedBasename-")), + s"expected every data file to be named '$expectedBasename-...', found: " + written.mkString(", ")) checkAnswer(spark.read.parquet(outputPath), df) } @@ -1350,8 +1341,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("dynamic partition overwrite falls back to Spark") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // A dynamic overwrite is a partitioned write, which CometWriteFiles declines - but the + // A dynamic overwrite is a partitioned write, which both native writers decline - but the // consequence of getting it wrong is silent data loss across untouched partitions, so assert // the fallback and the semantics explicitly rather than relying on the partitioning check. withTempPath { dir => @@ -1411,14 +1401,14 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("empty input still writes a schema-only file (SPARK-23271)") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // An empty input must still leave a schema behind for downstream readers: `spark.read.parquet` // of the output must see the write's schema, not fail. Comet reaches this in two ways - if the - // native child has one partition producing no batches, the partition-0 branch of executeTask - // writes a metadata-only file; if it produces zero partitions, doExecuteWrite swaps in a dummy - // single-partition RDD to get to the same branch. This test exercises the first; the - // zero-partition swap is reached by an AQE-collapsed empty relation and is covered by - // CometEmptyRelationParquetWriterSuite. + // native child has one partition producing no batches, the partition-0 branch of the write + // task writes a metadata-only file; if it produces zero partitions, the writer swaps in a + // dummy single-partition RDD to get to the same branch. Both CometWriteFilesExec (Spark 4.0+) + // and CometNativeWriteExec (Spark 3.x) do this. This test exercises the first; the + // zero-partition swap is covered by the empty-directory test below and, on Spark 4.0+, + // by CometEmptyRelationParquetWriterSuite. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val sourcePath = new File(dir, "source.parquet").getAbsolutePath @@ -1439,6 +1429,31 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a zero-partition native scan writes a schema-only file (#5303)") { + withTempPath { dir => + val source = new File(dir, "empty-source") + assert(source.mkdirs()) + val output = new File(dir, "output") + withNativeWriter { + val empty = spark.read.schema("id INT, name STRING").parquet(source.getAbsolutePath) + val scans = collect(empty.queryExecution.executedPlan) { case scan: CometNativeScanExec => + scan + } + assert(scans.size == 1, empty.queryExecution.executedPlan.toString) + assert(scans.head.executeColumnar().getNumPartitions == 0) + + val plan = captureWritePlan(p => empty.write.parquet(p), output.getAbsolutePath) + assertHasCometNativeWriteExec(plan) + assert(new File(output, "_SUCCESS").isFile) + assert(listPartFileNames(output.getAbsolutePath).size == 1) + assert(!new File(output, "_temporary").exists()) + val readBack = spark.read.parquet(output.getAbsolutePath) + assert(readBack.schema == empty.schema) + assert(readBack.collect().isEmpty) + } + } + } + test("INSERT INTO ... SELECT writes the target table's column names") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3426, which is Spark's own @@ -1487,8 +1502,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("an empty partition writes no file and still commits") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // executeTask's `sparkPartitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile + // The write task's `partitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile // altogether and still commit the task, matching FileFormatWriter's EmptyDirectoryDataWriter. // Hash-partitioning into eight and keeping a single id leaves at most one partition with // rows, so at most two files can appear: that one, plus partition 0's schema-only file when @@ -1547,8 +1561,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("a failing task aborts and cleans up its staging file") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // CometWriteFilesExec.executeTask must call committer.abortTask and rethrow. Injecting the + // Both native writers must call committer.abortTask and rethrow. Injecting the // failure through the commit protocol rather than the data lets the write get as far as // creating a staging file, so the cleanup is actually observable. // diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index d0ca67c60d5..40281ce4ea1 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala @@ -117,7 +117,10 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { command, child, outputPath = "/tmp/unused-native-write", - mode = SaveMode.Overwrite) + mode = SaveMode.Overwrite, + committer = null, + serializableHadoopConf = null, + outputWriterFactory = null) } private def assertRestoredParquetWrite(reverted: SparkPlan): WriteFilesExec = { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala new file mode 100644 index 00000000000..098b1c16bc5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -0,0 +1,390 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import java.io.{File, IOException} +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.jdk.CollectionConverters._ +import scala.util.Using + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.hadoop.mapreduce.{JobContext, TaskAttemptContext} +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.parquet.hadoop.{ParquetFileReader, ParquetOutputFormat} +import org.apache.parquet.hadoop.metadata.CompressionCodecName +import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.spark.TaskContext +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage +import org.apache.spark.internal.io.FileNameSpec +import org.apache.spark.sql.{CometTestBase, DataFrame, SaveMode} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, SQLHadoopMapReduceCommitProtocol} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{StringType, StructField} +import org.apache.spark.util.SerializableConfiguration + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.{Incompatible, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{schema2Proto, CometDataWritingCommand, NativeWriteUtils} + +/** Exercises the retained Spark 3.x writer, including both terminal execution entry points. */ +class CometNativeWriteSuite extends CometTestBase { + import NativeWriteCommitProtocol._ + + private def withWriter(partitions: Int = 1)( + f: (CometNativeWriteExec, DataFrame, File) => Unit): Unit = { + assume(!isSpark40Plus, "Spark 4.0+ uses CometWriteFilesExec") + withTempPath { dir => + val source = new File(dir, "source").getAbsolutePath + val output = new File(dir, "output with % space") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 20, 1, partitions).write.parquet(source) + } + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key -> classOf[NativeWriteCommitProtocol].getName, + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> "0") { + val data = spark.read.parquet(source) + val sourcePlan = data.queryExecution.executedPlan + val child = sourcePlan + .collectFirst { case scan: CometNativeScanExec => + scan + } + .getOrElse(fail(s"Expected a native source scan: $sourcePlan")) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path(output.toURI), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map("compression" -> "gzip", "comet.test.option" -> "present"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, child) + val nativeOp = CometDataWritingCommand.convert(command, Operator.newBuilder()).get + val writer = CometDataWritingCommand + .createExec(nativeOp, command) + .asInstanceOf[CometNativeWriteExec] + assert(writer.committer.isInstanceOf[NativeWriteCommitProtocol]) + reset() + try f(writer, data, output) + finally reset() + } + } + } + + test("Spark 3.x HDFS admission ignores the configurable output basename") { + assume(!isSpark40Plus, "Spark 4.0+ honors mapreduce.output.basename") + val data = spark.range(1) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path("hdfs://localhost:8020/output"), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map(NativeWriteUtils.BASE_OUTPUT_NAME -> "out%?#"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, data.queryExecution.executedPlan) + assert(CometDataWritingCommand.getSupportLevel(command).isInstanceOf[Incompatible]) + val escapedPath = cmd.copy(outputPath = new Path("hdfs://localhost:8020/out%?#")) + assert( + CometDataWritingCommand + .getSupportLevel(command.copy(cmd = escapedPath)) + .isInstanceOf[Unsupported]) + } + + Seq(false, true).foreach { columnar => + test( + s"configured commit protocol completes the ${if (columnar) "columnar" else "row"} write") { + withWriter(2) { (writer, data, output) => + val resultCount = + if (columnar) writer.executeColumnar().count() else writer.execute().count() + assert(resultCount == 0) + assert(events.asScala.head == "setupJob") + assert(events.asScala.last == "commitJob") + assert(!events.asScala.exists(_.startsWith("abort"))) + assert(writer.metrics("rows_written").value == 20, "the writer must execute natively") + val files = allocated.asScala.map(p => new Path(p).getName).toSet + assert(files.nonEmpty) + assert(files.forall(n => n.startsWith("chosen % ") && n.endsWith(".gz.parquet"))) + assert(output.list().filter(_.endsWith(".parquet")).toSet == files) + assert(new File(output, "_SUCCESS").isFile) + assert(!new File(output, "_temporary").exists()) + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + } + + test("the footer codec matches the codec in the file name") { + withWriter() { (writer, data, output) => + // The planner chose gzip, but the job configuration says zstd. The job configuration also + // names the file. + val conf = new Configuration(writer.serializableHadoopConf.value) + conf.set(ParquetOutputFormat.COMPRESSION, CompressionCodecName.ZSTD.name()) + writer.copy(serializableHadoopConf = new SerializableConfiguration(conf)).execute().count() + + val files = output.listFiles().filter(_.getName.endsWith(".parquet")) + assert(files.nonEmpty) + files.foreach { file => + assert(file.getName.endsWith(".zstd.parquet"), s"Expected a zstd file: ${file.getName}") + val input = HadoopInputFile.fromPath(new Path(file.toURI), conf) + Using.resource(ParquetFileReader.open(input)) { reader => + val codecs = reader.getFooter.getBlocks.asScala + .flatMap(_.getColumns.asScala) + .map(_.getCodec) + .toSet + assert(codecs == Set(CompressionCodecName.ZSTD), s"${file.getName} uses $codecs") + } + } + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + + test("native iterator creation failure aborts the task and job") { + withWriter() { (writer, _, output) => + // Inject an unknown native config after the child scan has initialized, so the writer's + // CometExecIterator constructor fails in createPlan, before the drain/close block is entered. + failAt = "createIterator" + failAbort = true + val error = intercept[Exception](writer.execute().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "createPlan"))) + assert(causes.exists(e => Option(e.getMessage).exists(_.contains("no_such_namespace")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("native execution failure survives throwing cleanup and abort callbacks") { + withWriter() { (writer, _, output) => + // A mismatched output type fails when writing the first native batch, after the native + // stream is initialized. This also lets releasePlan exercise its final metrics callback. + val badWriter = writer.nativeOp.getParquetWriter.toBuilder + .clearOutputSchema() + .addAllOutputSchema(schema2Proto(Seq(StructField("id", StringType))).asJava) + .build() + val badPlan = writer.nativeOp.toBuilder.setParquetWriter(badWriter).build() + failAbort = true + val failingWriter = new CleanupFailingNativeWriteExec(writer.copy(nativeOp = badPlan)) + val error = intercept[Exception](failingWriter.executeColumnar().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "executePlan"))) + assert(causes.exists(e => + Option(e.getMessage).exists(_.contains("Failed to rename batch columns")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected cleanup failure"))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("a cleanup failure aborts instead of committing") { + withWriter() { (writer, _, output) => + val error = + intercept[Exception](new CleanupFailingNativeWriteExec(writer).execute().count()) + assert(allErrors(error).exists(_.getMessage == "injected cleanup failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + Seq("commitTask", "onTaskCommit", "commitJob").foreach { phase => + test(s"$phase failure aborts and preserves the commit failure") { + withWriter() { (writer, _, output) => + failAt = phase + failAbort = true + val error = intercept[Exception](writer.execute().count()) + assert(allErrors(error).exists(_.getMessage == s"injected $phase failure")) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assert(events.contains("abortJob")) + if (phase == "commitTask") { + assert(events.contains("abortTask")) + assert( + allErrors(error).exists( + _.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + } + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert( + written.asScala.nonEmpty, + "the failure must occur after the native file was written") + } + } + } + + private def allErrors(error: Throwable): Seq[Throwable] = + Seq(error) ++ Option(error.getCause).toSeq.flatMap(allErrors) ++ + error.getSuppressed.toSeq.flatMap(allErrors) + + private def assertAborted(output: File): Unit = { + assert(events.contains("abortTask")) + assert(events.contains("abortJob")) + assert(!events.contains("commitJob")) + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert(allocated.asScala.nonEmpty, "the test must reach task file allocation") + } +} + +/** + * Local-mode probes are reset for each write; executor and driver callbacks share this object. + */ +object NativeWriteCommitProtocol { + val events = new ConcurrentLinkedQueue[String]() + val allocated = new ConcurrentLinkedQueue[String]() + val written = new ConcurrentLinkedQueue[String]() + val received = new ConcurrentLinkedQueue[TaskCommitMessage]() + @volatile var failAt: String = "" + @volatile var failAbort = false + + def reset(): Unit = { + events.clear() + allocated.clear() + written.clear() + received.clear() + failAt = "" + failAbort = false + } +} + +class NativeWriteCommitProtocol(jobId: String, path: String, dynamicPartitionOverwrite: Boolean) + extends SQLHadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite) { + import NativeWriteCommitProtocol._ + private var taskFile: String = _ + + override def setupJob(context: JobContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getOutputFormatClass == classOf[ParquetOutputFormat[_]]) + assert(FileOutputFormat.getOutputPath(context) == new Path(path)) + assert(context.getConfiguration.get("spark.sql.sources.writeJobUUID") != null) + super.setupJob(context) + context.getConfiguration.set("comet.test.setupJob", "present") + events.add("setupJob") + } + + override def setupTask(context: TaskAttemptContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getConfiguration.get("comet.test.setupJob") == "present") + super.setupTask(context) + events.add("setupTask") + } + + override def newTaskTempFile( + context: TaskAttemptContext, + dir: Option[String], + spec: FileNameSpec): String = { + val sparkPath = new Path(super.newTaskTempFile(context, dir, spec)) + taskFile = new Path(sparkPath.getParent, "chosen % " + sparkPath.getName).toString + allocated.add(taskFile) + if (failAt == "createIterator") { + val properties = TaskContext.get().getLocalProperties + properties.setProperty(CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key, "true") + properties.setProperty("spark.comet.datafusion.no_such_namespace.option", "1") + } + taskFile + } + + override def commitTask(context: TaskAttemptContext): TaskCommitMessage = { + val file = new Path(taskFile) + assert(file.getFileSystem(context.getConfiguration).getFileStatus(file).getLen > 0) + written.add(taskFile) + events.add("commitTask") + if (failAt == "commitTask") throw new IOException("injected commitTask failure") + new TaskCommitMessage((context.getTaskAttemptID.toString, super.commitTask(context))) + } + + override def onTaskCommit(message: TaskCommitMessage): Unit = { + events.add("onTaskCommit") + received.add(message) + if (failAt == "onTaskCommit") throw new IOException("injected onTaskCommit failure") + super.onTaskCommit(message.obj.asInstanceOf[(String, TaskCommitMessage)]._2) + } + + override def commitJob(context: JobContext, messages: Seq[TaskCommitMessage]): Unit = { + assert(messages.nonEmpty) + assert(messages.toSet == received.asScala.toSet) + assert( + messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._1).distinct.size == + messages.size) + events.add("commitJob") + if (failAt == "commitJob") throw new IOException("injected commitJob failure") + super.commitJob(context, messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._2)) + } + + override def abortTask(context: TaskAttemptContext): Unit = { + events.add("abortTask") + super.abortTask(context) + if (failAbort) throw new IOException("injected abortTask failure") + } + + override def abortJob(context: JobContext): Unit = { + events.add("abortJob") + super.abortJob(context) + if (failAbort) throw new IOException("injected abortJob failure") + } +} + +/** + * With periodic updates disabled, this metric throws only during native iterator teardown. + * CometMetricNode reports a sum metric with `add`, so both update methods throw. + */ +class CleanupFailingWriteMetric extends SQLMetric("sum") { + override def copy(): SQLMetric = new CleanupFailingWriteMetric + override def set(value: Long): Unit = throw new IOException("injected cleanup failure") + override def add(value: Long): Unit = throw new IOException("injected cleanup failure") +} + +class CleanupFailingNativeWriteExec(writer: CometNativeWriteExec) + extends CometNativeWriteExec( + writer.nativeOp, + writer.originalPlan, + writer.child, + writer.outputPath, + writer.mode, + writer.committer, + writer.serializableHadoopConf, + writer.outputWriterFactory, + writer.jobTrackerID) { + override lazy val metrics: Map[String, SQLMetric] = { + val failingMetric = new CleanupFailingWriteMetric + sparkContext.register(failingMetric) + writer.metrics.updated("rows_written", failingMetric) + } +} diff --git a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala index 04585612ab0..c739e420315 100644 --- a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala @@ -64,12 +64,8 @@ class CometEmptyRelationParquetWriterSuite extends CometParquetWriterTestBase { } if (nativeEmpty) { assert(collect(plan) { case e: CometEmptyRelationExec => e }.nonEmpty) - // A native empty relation is a zero-partition RDD, which is what - // CometDataWritingCommand declines on the Spark 3.x path: that writer only maps - // existing partitions, so no task runs and no file is written. Spark 4.0+ goes - // through the WriteFilesExec seam instead, where CometWriteFilesExec swaps in a - // dummy single-partition RDD exactly as Spark's own WriteFilesExec does, so - // partition 0 still writes the schema-only file the readback above needs. + // CometWriteFilesExec replaces the zero-partition native empty relation + // with one empty partition, so a task writes the schema-only file above. assertHasCometNativeWriteExec(plan) } else { // Spark's own EmptyRelationExec is not a Comet operator, so