diff --git a/.github/workflows/spark_sql_writer_tests.yml b/.github/workflows/spark_sql_writer_tests.yml index c813a8b39e1..ccc1f5445ce 100644 --- a/.github/workflows/spark_sql_writer_tests.yml +++ b/.github/workflows/spark_sql_writer_tests.yml @@ -117,6 +117,13 @@ jobs: spark-short-version: ${{ inputs.spark-version }} skip-native-build: true + - name: Pre-compile Spark SQL test classes + run: | + cd apache-spark + rm -rf /root/.m2/repository/org/apache/parquet + NOLINT_ON_COMPILE=true build/sbt -Dsbt.log.noformat=true -mem 3072 \ + 'sql/Test/compile' + - name: Run Parquet writer tests run: | cd apache-spark diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index dbbee713ae1..b98b2658ebd 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -218,18 +218,10 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Output file path (final destination) - output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol) - work_dir: String, - /// Job ID for tracking this write operation - job_id: Option, - /// Task attempt ID for this specific task - task_attempt_id: Option, + /// Full file path to write for this Spark task attempt + task_output_path: String, /// Compression codec compression: ParquetCompression, - /// Partition ID (from Spark TaskContext) - partition_id: i32, /// Column names to use in the output Parquet file column_names: Vec, /// Object store configuration options @@ -242,15 +234,10 @@ pub struct ParquetWriterExec { impl ParquetWriterExec { /// Create a new ParquetWriterExec - #[allow(clippy::too_many_arguments)] pub fn try_new( input: Arc, - output_path: String, - work_dir: String, - job_id: Option, - task_attempt_id: Option, + task_output_path: String, compression: ParquetCompression, - partition_id: i32, column_names: Vec, object_store_options: HashMap, ) -> Result { @@ -266,12 +253,8 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, - output_path, - work_dir, - job_id, - task_attempt_id, + task_output_path, compression, - partition_id, column_names, object_store_options, metrics: ExecutionPlanMetricsSet::new(), @@ -409,7 +392,7 @@ impl DisplayAs for ParquetWriterExec { write!( f, "ParquetWriterExec: path={}, compression={:?}", - self.output_path, self.compression + self.task_output_path, self.compression ) } DisplayFormatType::TreeRender => unimplemented!(), @@ -446,12 +429,8 @@ impl ExecutionPlan for ParquetWriterExec { match children.len() { 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.task_output_path.clone(), self.compression.clone(), - self.partition_id, self.column_names.clone(), self.object_store_options.clone(), )?)), @@ -476,8 +455,7 @@ 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 part_file = self.task_output_path.clone(); let compression = self.compression.to_parquet()?; let column_names = self.column_names.clone(); @@ -492,17 +470,6 @@ impl ExecutionPlan for ParquetWriterExec { .collect(); let output_schema = Arc::new(arrow::datatypes::Schema::new(fields)); - // Generate part file name for this partition - // If using FileCommitProtocol (work_dir is set), include task_attempt_id in the filename - let part_file = if let Some(attempt_id) = task_attempt_id { - format!( - "{}/part-{:05}-{:05}.parquet", - work_dir, self.partition_id, attempt_id - ) - } else { - format!("{}/part-{:05}.parquet", work_dir, self.partition_id) - }; - // Configure writer properties let props = WriterProperties::builder() .set_compression(compression) @@ -851,18 +818,15 @@ 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 task_output_path = + "hdfs://namenode:9000/user/test_parquet_writer_exec/part-00000-c000.parquet" + .to_string(); let column_names = vec!["id".to_string(), "name".to_string()]; let parquet_writer = ParquetWriterExec::try_new( memory_exec, - output_path, - work_dir, - None, // job_id - Some(123), // task_attempt_id + task_output_path, ParquetCompression::None, - 0, // partition_id column_names, HashMap::new(), // object_store_options )?; diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index c179c3b57c5..0707cb2d2cc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1853,18 +1853,15 @@ impl PhysicalPlanner { .map(|(k, v)| (k.clone(), v.clone())) .collect(); + let task_output_path = writer + .task_output_path + .clone() + .expect("task_output_path is provided"); + let parquet_writer = Arc::new(ParquetWriterExec::try_new( Arc::clone(&child.native_plan), - writer.output_path.clone(), - writer - .work_dir - .as_ref() - .expect("work_dir is provided") - .clone(), - writer.job_id.clone(), - writer.task_attempt_id, + task_output_path, codec, - self.partition, writer.column_names.clone(), object_store_options, )?); diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index ced87262f32..0b8eed67555 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -458,16 +458,8 @@ message ShuffleWriter { } message ParquetWriter { - string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; - // Working directory for temporary files (used by FileCommitProtocol) - // If not set, files are written directly to output_path - 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; // 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. @@ -475,6 +467,9 @@ message ParquetWriter { // configuration value "spark.hadoop.fs.s3a.access.key" will be stored as "fs.s3a.access.key" in // the map. map object_store_options = 8; + // Full temporary output file path returned by Spark's FileCommitProtocol for this task. + // The native writer must write exactly to this path so Spark can commit or abort it. + optional string task_output_path = 9; } enum AggregateMode { 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 8157f286825..c33bc6b9e95 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 @@ -20,17 +20,23 @@ package org.apache.comet.serde.operator import java.net.URI -import java.util.Locale +import java.util.{Locale, UUID} import scala.jdk.CollectionConverters._ +import org.apache.hadoop.fs.Path +import org.apache.hadoop.mapreduce.Job +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat import org.apache.parquet.hadoop.ParquetOutputFormat -import org.apache.spark.SparkException +import org.apache.spark.internal.io.FileCommitProtocol +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec} import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.util.SerializableConfiguration import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -132,11 +138,10 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec val writerOpBuilder = OperatorOuterClass.ParquetWriter .newBuilder() - .setOutputPath(outputPath) .setCompression(codec) .addAllColumnNames(cmd.query.output.map(_.name).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 + // The task_output_path is filled in by CometNativeWriteExec at execution time from + // Spark's FileCommitProtocol, because it depends on the Spark task context. // Collect S3/cloud storage configurations val session = op.session @@ -181,29 +186,34 @@ 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 hadoopConf = session.sessionState.newHadoopConfWithOptions(cmd.options) + val job = Job.getInstance(hadoopConf) + job.setOutputKeyClass(classOf[Void]) + job.setOutputValueClass(classOf[InternalRow]) + FileOutputFormat.setOutputPath(job, new Path(outputPath)) + + val outputWriterFactory = + cmd.fileFormat.prepareWrite(session, job, CaseInsensitiveMap(cmd.options), cmd.query.schema) + + val commitProtocolJobId = UUID.randomUUID().toString + val committer = FileCommitProtocol.instantiate( + session.sessionState.conf.fileCommitProtocolClass, + commitProtocolJobId, + outputPath, + false) + + // Match Spark's FileFormatWriter behavior: propagate a per-write UUID in the Hadoop + // configuration before it is serialized to executors. + job.getConfiguration.set("spark.sql.sources.writeJobUUID", UUID.randomUUID().toString) + + val commitProtocol = CometNativeWriteExec.CommitProtocolConfig( + committer = committer, + serializableHadoopConf = new SerializableConfiguration(job.getConfiguration), + outputWriterFactory = outputWriterFactory, + jobTrackerID = CometNativeWriteExec.newJobTrackerID()) + + CometNativeWriteExec(nativeOp, childPlan, outputPath, cmd.mode, commitProtocol) } private def parseCompressionCodec(cmd: InsertIntoHadoopFsRelationCommand) = { 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 f0d10b17667..3bd6947699e 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,79 +19,82 @@ 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 +import org.apache.spark.sql.comet.CometNativeWriteExec.CommitProtocolConfig 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 import com.google.protobuf.CodedOutputStream import org.apache.comet.CometExecIterator import org.apache.comet.serde.OperatorOuterClass.Operator +object CometNativeWriteExec { + + def newJobTrackerID(): String = SparkHadoopWriterUtils.createJobTrackerID(new Date()) + + /** + * Driver-created objects required to use Spark's FileCommitProtocol from native write tasks. + * The committer instance is serializable by contract and is sent to executors, while the same + * driver-side instance receives task commit messages and commits or aborts the job. + */ + case class CommitProtocolConfig( + committer: FileCommitProtocol, + serializableHadoopConf: SerializableConfiguration, + outputWriterFactory: OutputWriterFactory, + jobTrackerID: String) + extends Serializable +} + /** - * 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. + * Comet physical operator for native Parquet write operations. * - * 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. + * This operator follows Spark's FileCommitProtocol lifecycle: driver setupJob, executor + * setupTask/newTaskTempFile/commitTask or abortTask, and driver commitJob or abortJob. Native + * code writes exactly to the task temp file returned by Spark's commit protocol; it does not + * generate its own final part-file names. * * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) + * The native operator representing the write operation (template, modified per task) * @param child * The child operator providing the data to write * @param outputPath - * The path where the Parquet file will be written + * The final output directory for the write * @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 + * The Spark SaveMode governing target-exists behavior + * @param commitProtocol + * Spark file commit protocol state */ case class CometNativeWriteExec( nativeOp: Operator, child: SparkPlan, outputPath: String, mode: SaveMode, - committer: Option[FileCommitProtocol] = None, - jobTrackerID: String = Utils.createTempDir().getName) + commitProtocol: CommitProtocolConfig) 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(serializeNativeOp(nativeOp))) override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) @@ -104,64 +107,69 @@ 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() + // Write operations do not return rows. 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() + // Write operations do not return columnar batches. Spark may still ask for columnar output + // while this operator is nested under write planning nodes, so run the terminal write here too. + sparkContext.emptyRDD[ColumnarBatch] + } + + private def executeWriteAndCommit(): Unit = { if (!prepareOutputPathForMode()) { logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") - return sparkContext.emptyRDD[ColumnarBatch] + return + } + + val jobContext = createJobContext(commitProtocol) + + // Match Spark's FileFormatWriter lifecycle: setupJob is outside the try block because it + // only initializes the job; failures after this point should abort the job. + commitProtocol.committer.setupJob(jobContext) + + try { + val commitMessages = runNativeWriteJob(commitProtocol) + commitProtocol.committer.commitJob(jobContext, commitMessages.toSeq) + + val filesWritten = metrics("files_written").value + val bytesWritten = metrics("bytes_written").value + val rowsWritten = metrics("rows_written").value + logInfo( + s"Successfully committed native write job to $outputPath: " + + s"$filesWritten files, $bytesWritten bytes, $rowsWritten rows") + } catch { + case t: Throwable => + abortJob(commitProtocol, jobContext, t) + throw t } + } - // Get the input data from the child operator + private def runNativeWriteJob(protocol: CommitProtocolConfig): Array[TaskCommitMessage] = { + val writeRDD = nativeWriteTasks(protocol) + val ret = new Array[TaskCommitMessage](writeRDD.partitions.length) + + sparkContext.runJob( + writeRDD, + (_: TaskContext, iter: Iterator[TaskCommitMessage]) => { + assert(iter.hasNext, "Native write task did not return a commit message") + val commitMessage = iter.next() + assert(!iter.hasNext, "Native write task returned more than one commit message") + commitMessage + }, + writeRDD.partitions.indices, + (index, commitMessage: TaskCommitMessage) => { + protocol.committer.onTaskCommit(commitMessage) + ret(index) = commitMessage + }) + + ret + } + + private def nativeWriteTasks(protocol: CommitProtocolConfig): RDD[TaskCommitMessage] = { val childRDD = if (child.supportsColumnar) { child.executeColumnar() } else { @@ -174,137 +182,94 @@ case class CometNativeWriteExec( } } - // Capture metadata before the transformation val numPartitions = childRDD.getNumPartitions val numOutputCols = child.output.length - val capturedCommitter = committer - val capturedJobTrackerID = jobTrackerID val capturedNativeOp = nativeOp - val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks + val writeExec = this - // 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, commitMsg) = capturedCommitter - .map { committer => - val taskContext = - createTaskContext(capturedJobTrackerID, partitionId, taskAttemptId.toInt) - - // Setup task - this creates the temporary working directory - committer.setupTask(taskContext) + val sparkTaskContext = TaskContext.get() + val partitionId = sparkTaskContext.partitionId() + val sparkStageId = sparkTaskContext.stageId() + val taskAttemptId = sparkTaskContext.taskAttemptId() + val sparkAttemptNumber = taskAttemptId.toInt & Integer.MAX_VALUE - // 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 - } + new Iterator[TaskCommitMessage] { + private var emitted = false - 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 - } + override def hasNext: Boolean = !emitted - if (!result && !completed) { - handleTaskEnd() + override def next(): TaskCommitMessage = { + if (emitted) { + throw new NoSuchElementException("Native write task already completed") } + emitted = true - result - } + val taskAttemptContext = + createTaskContext(protocol, sparkStageId, partitionId, sparkAttemptNumber) + protocol.committer.setupTask(taskAttemptContext) - override def next(): ColumnarBatch = { - try { - execIterator.next() - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - } + var execIterator: CometExecIterator = null - private def handleTaskEnd(): Unit = { - if (!completed) { - completed = true + try { + val fileExtension = protocol.outputWriterFactory.getFileExtension(taskAttemptContext) + val taskOutputPath = protocol.committer.newTaskTempFile( + taskAttemptContext, + None, + FileNameSpec("", "-c000" + fileExtension)) + + val parquetWriter = capturedNativeOp.getParquetWriter.toBuilder + .setTaskOutputPath(taskOutputPath) + .build() + + val modifiedNativeOp = capturedNativeOp.toBuilder + .setParquetWriter(parquetWriter) + .build() + + val nativeMetrics = CometMetricNode.fromCometPlan(writeExec) + // 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) + + execIterator = new CometExecIterator( + CometExec.newIterId, + CometArrowStream.inputObjects( + iter, + CometUtils.fromAttributes(child.output), + "CometNativeWriteExec"), + numOutputCols, + serializeNativeOp(modifiedNativeOp), + nativeMetrics, + numPartitions, + partitionId, + None, + Seq.empty) + + while (execIterator.hasNext) { + execIterator.next().close() + } - // Handle commit or abort based on whether an exception was thrown - taskContext.foreach { case (committer, ctx) => + val message = protocol.committer.commitTask(taskAttemptContext) + logInfo(s"Task ${taskAttemptContext.getTaskAttemptID} committed successfully") + message + } catch { + case t: Throwable => 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") - } + protocol.committer.abortTask(taskAttemptContext) + logWarning( + s"Task ${taskAttemptContext.getTaskAttemptID} aborted due to exception: " + + Option(t.getMessage).getOrElse(t.getClass.getName)) } 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 - } + case abortError: Throwable => + logWarning( + s"Error aborting task ${taskAttemptContext.getTaskAttemptID}", + abortError) + t.addSuppressed(abortError) } + throw t + } finally { + if (execIterator != null) { + execIterator.close() } } } @@ -312,11 +277,32 @@ case class CometNativeWriteExec( } } - /** 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 + private def serializeNativeOp(op: Operator): Array[Byte] = { + val size = op.getSerializedSize + val bytes = new Array[Byte](size) + val codedOutput = CodedOutputStream.newInstance(bytes) + op.writeTo(codedOutput) + codedOutput.checkNoSpaceLeft() + bytes + } + + private def abortJob( + protocol: CommitProtocolConfig, + jobContext: Job, + cause: Throwable): Unit = { + logError("Native write failed, aborting job", cause) + try { + protocol.committer.abortJob(jobContext) + } catch { + case abortError: Throwable => + logWarning("Error aborting native write job", abortError) + cause.addSuppressed(abortError) + } + } + + /** Create a JobContext for the write job using the prepared Hadoop write configuration. */ + private def createJobContext(protocol: CommitProtocolConfig): Job = { + Job.getInstance(new Configuration(protocol.serializableHadoopConf.value)) } /** @@ -329,7 +315,7 @@ case class CometNativeWriteExec( */ private def prepareOutputPathForMode(): Boolean = { val path = new Path(outputPath) - val hadoopConf = sparkContext.hadoopConfiguration + val hadoopConf = commitProtocol.serializableHadoopConf.value val fs = path.getFileSystem(hadoopConf) val qualifiedOutputPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) @@ -343,11 +329,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 (!commitProtocol.committer.deleteWithJob(fs, qualifiedOutputPath, true)) { throw QueryExecutionErrors.cannotClearOutputDirectoryError(qualifiedOutputPath) } } @@ -359,13 +341,21 @@ 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) + protocol: CommitProtocolConfig, + sparkStageId: Int, + sparkPartitionId: Int, + sparkAttemptNumber: Int): TaskAttemptContext = { + val hadoopConf = new Configuration(protocol.serializableHadoopConf.value) + val jobId = SparkHadoopWriterUtils.createJobID(protocol.jobTrackerID, sparkStageId) + val taskId = new TaskID(jobId, TaskType.MAP, sparkPartitionId) + val taskAttemptId = new TaskAttemptID(taskId, sparkAttemptNumber) + + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskAttemptId.getTaskID.toString) + hadoopConf.set("mapreduce.task.attempt.id", taskAttemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + + new TaskAttemptContextImpl(hadoopConf, taskAttemptId) } } 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 a1ae1af1d1c..d9398b4421c 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -81,6 +81,7 @@ class CometParquetWriterSuite extends CometTestBase { writeWithCometNativeWriteExec(inputPath, outputPath) verifyWrittenFile(outputPath) + verifyCommitProtocolOutput(outputPath) } } } @@ -801,6 +802,28 @@ class CometParquetWriterSuite extends CometTestBase { Some(plan) } + private def verifyCommitProtocolOutput(outputPath: String): Unit = { + val outputDir = new File(outputPath) + val outputFiles = Option(outputDir.listFiles()).getOrElse(Array.empty) + val fileNames = outputFiles.map(_.getName).toSeq + + assert( + fileNames.contains("_SUCCESS"), + s"Expected Spark commit protocol to create _SUCCESS marker, found: ${fileNames.mkString(", ")}") + assert( + !fileNames.contains("_temporary"), + s"Expected temporary commit directory to be cleaned up, found: ${fileNames.mkString(", ")}") + assert( + !fileNames.exists(_.startsWith(".spark-staging-")), + s"Expected staging commit directory to be cleaned up, found: ${fileNames.mkString(", ")}") + + val partFileNames = fileNames.filter(_.startsWith("part-")) + assert(partFileNames.nonEmpty, s"Expected part files, found: ${fileNames.mkString(", ")}") + assert( + partFileNames.forall(name => name.contains("-c000") && name.endsWith(".parquet")), + s"Expected Spark commit protocol part-file names, found: ${partFileNames.mkString(", ")}") + } + private def verifyWrittenFile(outputPath: String): Unit = { // Verify the data was written correctly val resultDf = spark.read.parquet(outputPath)