fix: honor Spark’s commit protocol in Spark 3.x native writes

.github/workflows/pr_build_linux.yml

@@ -472,6 +472,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

.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

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

@@ -220,18 +220,9 @@ impl ParquetWriter {
 pub struct ParquetWriterExec {
     /// Input execution plan
     input: Arc<dyn ExecutionPlan>,
-    /// 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<String>,
-    /// Job ID for tracking this write operation
-    job_id: Option<String>,
-    /// Task attempt ID for this specific task
-    task_attempt_id: Option<i32>,
     /// Compression codec
     compression: ParquetCompression,
     /// Partition ID (from Spark TaskContext)
@@ -254,9 +245,6 @@ impl ParquetWriterExec {
     pub fn try_new(
         input: Arc<dyn ExecutionPlan>,
         output_path: String,
-        work_dir: Option<String>,
-        job_id: Option<String>,
-        task_attempt_id: Option<i32>,
         compression: ParquetCompression,
         partition_id: i32,
         column_names: Vec<String>,
@@ -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,

native/core/src/execution/planner.rs

@@ -1998,9 +1998,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(),

native/proto/src/proto/operator.proto

@@ -897,27 +897,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.

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-<partition>-<attempt>.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.
 
       // 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)
   }
 
 }

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 `<basename>-<split>-<jobId>`; 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 `<basename>-<split>-<jobId>`; 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