feat(parquet): implement file commit protocol for native Parquet writes

.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

native/core/src/execution/operators/parquet_writer.rs

@@ -218,18 +218,10 @@ impl ParquetWriter {
 pub struct ParquetWriterExec {
     /// Input execution plan
     input: Arc<dyn ExecutionPlan>,
-    /// 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<String>,
-    /// Task attempt ID for this specific task
-    task_attempt_id: Option<i32>,
+    /// 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<String>,
     /// 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<dyn ExecutionPlan>,
-        output_path: String,
-        work_dir: String,
-        job_id: Option<String>,
-        task_attempt_id: Option<i32>,
+        task_output_path: String,
         compression: ParquetCompression,
-        partition_id: i32,
         column_names: Vec<String>,
         object_store_options: HashMap<String, String>,
     ) -> Result<Self> {
@@ -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
         )?;

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,
                 )?);

native/proto/src/proto/operator.proto

@@ -458,23 +458,18 @@ 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.
   // The configuration values have hadoop. or spark.hadoop. prefix trimmed. For instance, the
   // configuration value "spark.hadoop.fs.s3a.access.key" will be stored as "fs.s3a.access.key" in
   // the map.
   map<string, string> 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 {

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) = {