From 0a3a98b69fd7bfd375bd7a045832f7874092d823 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 00:21:16 +0800 Subject: [PATCH 01/10] =?UTF-8?q?fix:=20honor=20Spark=E2=80=99s=20commit?= =?UTF-8?q?=20protocol=20in=20Spark=203.x=20native=20writes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../operator/CometDataWritingCommand.scala | 66 ++-- .../sql/comet/CometNativeWriteExec.scala | 359 ++++++------------ .../parquet/CometParquetWriterSuite.scala | 13 +- .../sql/comet/CometNativeWriteSuite.scala | 329 ++++++++++++++++ 4 files changed, 478 insertions(+), 289 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala 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 6af2cc18ba6..acedd25a589 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,9 +19,15 @@ package org.apache.comet.serde.operator +import java.util.UUID + import scala.jdk.CollectionConverters._ -import org.apache.spark.SparkException +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.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.QueryStageExec @@ -30,6 +36,7 @@ import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCom 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 @@ -70,10 +77,11 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec return Unsupported(Some("Supported output filesystems: local, HDFS")) } + val hadoopConf = op.session.sessionState.newHadoopConfWithOptions(cmd.options) 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, + hadoopConf.get(NativeWriteUtils.BASE_OUTPUT_NAME, "part")) .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -140,8 +148,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, leaving work_dir unset. // Collect S3/cloud storage configurations val session = op.session @@ -189,29 +197,29 @@ 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, 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, + childPlan, + outputPath, + cmd.mode, + committer, + new SerializableConfiguration(job.getConfiguration), + outputWriterFactory) } } 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 ced525c1869..e4f076eaed6 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,15 @@ 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.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 @@ -33,65 +35,37 @@ import org.apache.spark.sql.comet.execution.arrow.CometArrowStream 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.datasources.OutputWriterFactory 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. - * - * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) - * @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 + * 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. */ case class CometNativeWriteExec( nativeOp: Operator, 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 { override def originalPlan: SparkPlan = child - // 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) @@ -104,219 +78,104 @@ 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 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 = 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 - val capturedCommitter = committer - val capturedJobTrackerID = jobTrackerID + val childSchema = CometUtils.fromAttributes(child.output) 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 capturedCommitter = committer + val writerFactory = outputWriterFactory + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val commitMessages = new Array[TaskCommitMessage](numPartitions) + + sparkContext.runJob( + childRDD, + (context: TaskContext, batches: Iterator[ColumnarBatch]) => { + val taskContext = createTaskContext(hadoopConf.value, 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 = { + val extension = writerFactory.getFileExtension(taskContext) + val filePath = capturedCommitter.newTaskTempFile( + taskContext, + None, + FileNameSpec("", "-c000" + extension)) + NativeWriteUtils.checkNativeWriteDestination(filePath) + val writer = capturedNativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .clearWorkDir() + .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 + capturedCommitter.commitTask(taskContext) + })(catchBlock = capturedCommitter.abortTask(taskContext)) + }, + childRDD.partitions.indices, + (index, message: TaskCommitMessage) => { + committer.onTaskCommit(message) + commitMessages(index) = message + }) + commitMessages } /** @@ -329,7 +188,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) @@ -343,11 +202,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) } } @@ -357,15 +212,17 @@ case class CometNativeWriteExec( } } - /** Create a TaskAttemptContext for a specific task */ - 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) + /** Match FileFormatWriter's Hadoop task identifiers and configuration. */ + private def createTaskContext(conf: Configuration, 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/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index babe86b3c72..68e0888829b 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -1080,16 +1080,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 => @@ -1109,7 +1106,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. @@ -1390,8 +1386,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/spark/sql/comet/CometNativeWriteSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala new file mode 100644 index 00000000000..ad9eec00c83 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -0,0 +1,329 @@ +/* + * 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 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.ParquetOutputFormat +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.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{schema2Proto, CometDataWritingCommand} + +/** 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() + } + } + } + + 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("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. */ +class CleanupFailingWriteMetric extends SQLMetric("sum") { + override def copy(): SQLMetric = new CleanupFailingWriteMetric + override def set(value: Long): Unit = throw new IOException("injected cleanup failure") +} + +class CleanupFailingNativeWriteExec(writer: CometNativeWriteExec) + extends CometNativeWriteExec( + writer.nativeOp, + 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) + } +} From d2f4ef1065e342a19206bf90ff30458459376dda Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 15 Sep 2026 00:35:38 +0800 Subject: [PATCH 02/10] ci: register legacy native write commit protocol tests --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index ff12e42d497..d49e8640b1a 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -503,6 +503,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.ParquetReadV2Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index b01ec62b006..448fea05686 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -120,6 +120,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.ParquetReadV2Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite From b9b9e0d39ad1b60649ff5cc60755c75550b73a62 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 25 Sep 2026 19:16:02 +0800 Subject: [PATCH 03/10] refactor: drop the native Parquet writer's work_dir file naming Both native writers now take the exact file name from Spark's commit protocol, so nothing sets work_dir, job_id or task_attempt_id any more. Remove the Some(work_dir) arm from ParquetWriterExec::execute, reserve the three proto fields, and correct the comments that still described the Spark 3.x writer naming its own files. --- .../src/execution/operators/parquet_writer.rs | 69 ++++--------------- native/core/src/execution/planner.rs | 3 - native/proto/src/proto/operator.proto | 28 +++----- .../operator/CometDataWritingCommand.scala | 2 +- .../serde/operator/NativeWriteUtils.scala | 10 +-- 5 files changed, 29 insertions(+), 83 deletions(-) diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 9897ed626c6..96dc9be0b9f 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -220,18 +220,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) @@ -254,9 +245,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, @@ -276,9 +264,6 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, output_path, - work_dir, - job_id, - task_attempt_id, compression, partition_id, column_names, @@ -464,9 +449,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(), @@ -494,8 +476,6 @@ impl ExecutionPlan for ParquetWriterExec { let runtime_env = context.runtime_env(); 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(); @@ -514,19 +494,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() @@ -654,11 +623,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), @@ -675,9 +643,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, @@ -742,13 +707,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()], @@ -811,13 +773,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()], @@ -1068,16 +1027,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.rs b/native/core/src/execution/planner.rs index 9e1f423f9b5..0ca2d00daed 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2050,9 +2050,6 @@ impl PhysicalPlanner { 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, self.partition, writer.column_names.clone(), diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 7b3a7aec0d2..09dc987e14b 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -900,27 +900,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 acedd25a589..958576b6aba 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 @@ -149,7 +149,7 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec Some( op.session.sessionState.conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))).asJava) // CometNativeWriteExec replaces output_path with the committer's exact task filename - // at execution time, leaving work_dir unset. + // at execution time. // Collect S3/cloud storage configurations val session = op.session 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 d8247e99a7f..942077ddddf 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,12 @@ 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 + * - `fileNamePrefix`, the basename every file name is built from. 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. + * interpolates into `--`; both native writers take their file names + * from the commit protocol, on every supported Spark version. 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. * * 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 From b1dba50bb0cf880e00af99c149374ae50f6b9cc0 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 25 Sep 2026 19:16:02 +0800 Subject: [PATCH 04/10] fix: handle empty inputs in the Spark 3.x native writer Mirror CometWriteFilesExec and FileFormatWriter: swap a zero-partition child for a single empty partition so the output still carries a schema (SPARK-23271), and only write a file from partition 0 or a partition that has rows. Move createTaskContext to the companion object so the task closure no longer captures the exec node and its child subtree. The SPARK-23271 and empty-partition tests now run on Spark 3.x too. --- .../sql/comet/CometNativeWriteExec.scala | 104 +++++++++++------- .../parquet/CometParquetWriterSuite.scala | 11 +- 2 files changed, 72 insertions(+), 43 deletions(-) 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 e4f076eaed6..d16ad508723 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 @@ -118,59 +118,79 @@ case class CometNativeWriteExec( } } - val numPartitions = childRDD.getNumPartitions + // 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 writerFactory = outputWriterFactory val nativeMetrics = CometMetricNode.fromCometPlan(this) val commitMessages = new Array[TaskCommitMessage](numPartitions) sparkContext.runJob( - childRDD, + writeRDD, (context: TaskContext, batches: Iterator[ColumnarBatch]) => { - val taskContext = createTaskContext(hadoopConf.value, context) + 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 = { - val extension = writerFactory.getFileExtension(taskContext) - val filePath = capturedCommitter.newTaskTempFile( - taskContext, - None, - FileNameSpec("", "-c000" + extension)) - NativeWriteUtils.checkNativeWriteDestination(filePath) - val writer = capturedNativeOp.getParquetWriter.toBuilder - .setOutputPath(filePath) - .clearWorkDir() - .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() + // 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) + val writer = capturedNativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .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() } - } { - execIterator.close() } + // `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)) }, - childRDD.partitions.indices, + writeRDD.partitions.indices, (index, message: TaskCommitMessage) => { committer.onTaskCommit(message) commitMessages(index) = message @@ -212,8 +232,18 @@ case class CometNativeWriteExec( } } - /** Match FileFormatWriter's Hadoop task identifiers and configuration. */ - private def createTaskContext(conf: Configuration, context: TaskContext): TaskAttemptContext = { +} + +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( + 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()) 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 68e0888829b..28c5e4a432f 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -1250,12 +1250,12 @@ 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 + // 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 reached by an AQE-collapsed empty relation and is covered by // CometEmptyRelationParquetWriterSuite. withTempPath { dir => @@ -1326,8 +1326,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 From c5868dcdef57ccea4511738019bbf6874155bb6a Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 29 Sep 2026 00:47:22 +0800 Subject: [PATCH 05/10] fix: match Spark 3.x write admission and cover empty scans --- .../operator/CometDataWritingCommand.scala | 29 ++---------- .../serde/operator/NativeWriteUtils.scala | 13 +++--- .../sql/comet/CometNativeWriteExec.scala | 2 +- .../parquet/CometParquetWriterSuite.scala | 46 ++++++++++++++----- .../sql/comet/CometNativeWriteSuite.scala | 28 ++++++++++- ...CometEmptyRelationParquetWriterSuite.scala | 8 +--- 6 files changed, 76 insertions(+), 50 deletions(-) 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 958576b6aba..8225f4b8121 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 @@ -28,13 +28,10 @@ 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.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec +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 @@ -62,26 +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")) } - val hadoopConf = op.session.sessionState.newHadoopConfWithOptions(cmd.options) + // 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 .escapedHdfsDestination( cmd.outputPath.toString, - hadoopConf.get(NativeWriteUtils.BASE_OUTPUT_NAME, "part")) + NativeWriteUtils.DEFAULT_BASE_OUTPUT_NAME) .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -106,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, 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 942077ddddf..05a9bce8bf8 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. That is - * `mapreduce.output.basename`, which `HadoopMapReduceCommitProtocol.getFilename` - * interpolates into `--`; both native writers take their file names - * from the commit protocol, on every supported Spark version. 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 d16ad508723..13eeb9957c9 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 @@ -109,7 +109,7 @@ case class CometNativeWriteExec( private def runNativeWriteJob( hadoopConf: SerializableConfiguration): Array[TaskCommitMessage] = { - val childRDD = if (child.supportsColumnar) { + val childRDD: RDD[ColumnarBatch] = if (child.supportsColumnar) { child.executeColumnar() } else { child.execute().mapPartitionsInternal { _ => 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 28c5e4a432f..e04bc782856 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -1134,12 +1134,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 @@ -1155,9 +1152,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) } @@ -1189,8 +1187,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 => @@ -1256,8 +1253,8 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { // 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 reached by an AQE-collapsed empty relation and is covered by - // CometEmptyRelationParquetWriterSuite. + // 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 @@ -1278,6 +1275,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 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 index ad9eec00c83..df84052747b 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -41,8 +41,9 @@ import org.apache.spark.sql.types.{StringType, StructField} 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} +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 { @@ -96,6 +97,31 @@ class CometNativeWriteSuite extends CometTestBase { } } + 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") { 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 18f9ee2ec0a..03c0ca7c150 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 From 19ef36062b99d9d5e0c47698bb6dba3ea3451751 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 2 Oct 2026 22:55:15 +0800 Subject: [PATCH 06/10] docs: say native writes over an empty relation stay native With the empty-relation guard gone from CometDataWritingCommand, the user guide and the CometWriteFilesExec comment still described Parquet writes over an empty relation falling back to Spark's writer. Both native write paths now swap a zero-partition input for one empty partition, so such a write stays native and still leaves a readable schema-only file. --- docs/source/user-guide/latest/compatibility/operators.md | 6 ++++-- docs/source/user-guide/latest/operators.md | 2 +- .../org/apache/spark/sql/comet/CometWriteFilesExec.scala | 8 ++++---- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 53b290726aa..4b7358f1f7f 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 c9ba5e491e1..97760abf964 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -54,7 +54,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. | | `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). | ## Projection and filtering 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 { From ce8a8ec074fd418076676c8e34ed037920af3cb8 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 4 Oct 2026 23:34:52 +0800 Subject: [PATCH 07/10] fix: take the Spark 3.x native write codec from the job configuration The Spark 3.x task named its file after the codec in the job configuration, because getFileExtension reads it from there. The native writer used the codec that convert resolved at planning time. The two agree today only because both copy the precedence rule of ParquetOptions. Read the codec with CodecConfig.from(taskContext) next to setOutputPath, as CometWriteFilesExec.executeTask does, so the file name and the footer agree by construction. A new CometNativeWriteSuite test sets a job codec that differs from the planned one and checks the name and the footer. The mixed-case codec test now checks the file name on Spark 3.x too. --- .../operator/CometDataWritingCommand.scala | 2 ++ .../sql/comet/CometNativeWriteExec.scala | 9 ++++++ .../parquet/CometParquetWriterSuite.scala | 11 +++---- .../sql/comet/CometNativeWriteSuite.scala | 32 ++++++++++++++++++- 4 files changed, 46 insertions(+), 8 deletions(-) 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 8225f4b8121..1831e2c1739 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 @@ -112,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 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 13eeb9957c9..641ee538e77 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 @@ -25,6 +25,7 @@ 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, SparkHadoopWriterUtils} import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage @@ -158,8 +159,16 @@ case class CometNativeWriteExec( 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() 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 e04bc782856..6e535c8df75 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -456,13 +456,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'") } } } 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 index df84052747b..69ee07261c5 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -23,11 +23,15 @@ 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.ParquetOutputFormat +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 @@ -38,6 +42,7 @@ 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 @@ -144,6 +149,31 @@ class CometNativeWriteSuite extends CometTestBase { } } + 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 From f4e639c50609d67f4684a1ab45bbf471e6614a57 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Mon, 5 Oct 2026 11:48:40 +0800 Subject: [PATCH 08/10] test: make the injected cleanup failure throw from SQLMetric.add too Since #6416, CometMetricNode reports a sum metric with add instead of set. The test metric threw only from set, so the two cleanup-failure tests in CometNativeWriteSuite stopped injecting their failure and failed on Spark 3.4 and 3.5. --- .../org/apache/spark/sql/comet/CometNativeWriteSuite.scala | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 index 69ee07261c5..6d0d9b21753 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -361,10 +361,14 @@ class NativeWriteCommitProtocol(jobId: String, path: String, dynamicPartitionOve } } -/** With periodic updates disabled, this metric throws only during native iterator teardown. */ +/** + * 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) From b09a70b9fa791e75ce04aee484508ac29039b0ab Mon Sep 17 00:00:00 2001 From: peterxcli Date: Tue, 14 Jul 2026 05:16:42 +0800 Subject: [PATCH 09/10] ci: pre-compile Spark SQL test classes in the manual writer workflow The Parquet writer test step caps the sbt heap at 1 GB for the forked test JVM, but this workflow compiles Spark first, and sbt ran out of heap while it compiled sql/core. Compile the SQL test classes in an earlier step with a 3 GB heap. This branch had the same step in July, before a history rewrite dropped it. --- .github/workflows/spark_sql_writer_tests.yml | 7 +++++++ 1 file changed, 7 insertions(+) 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 From 25afd98d774834a1a6ab3d125eb34d9360c0af52 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 7 Oct 2026 10:04:00 -0600 Subject: [PATCH 10/10] test: build CometNativeWriteExec with the merged constructor The merge with main brings in two callers written against the old constructor. #5957's RevertNativeForTransitionHeavyStagesSuite relied on the removed committer and jobTrackerID defaults, and #6247's write_reserving_from helper still passed the removed work_dir, job_id and task_attempt_id arguments to ParquetWriterExec::try_new. CleanupFailingNativeWriteExec also needs #5957's originalPlan argument. --- native/core/src/execution/operators/parquet_writer.rs | 3 --- .../rules/RevertNativeForTransitionHeavyStagesSuite.scala | 5 ++++- .../org/apache/spark/sql/comet/CometNativeWriteSuite.scala | 1 + 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index ed62202b6eb..cc2da616e54 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -862,9 +862,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, format!("file://{}", file.display()), - None, - None, - None, ParquetCompression::None, 0, column_names, 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 index 6d0d9b21753..098b1c16bc5 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -374,6 +374,7 @@ class CleanupFailingWriteMetric extends SQLMetric("sum") { class CleanupFailingNativeWriteExec(writer: CometNativeWriteExec) extends CometNativeWriteExec( writer.nativeOp, + writer.originalPlan, writer.child, writer.outputPath, writer.mode,