refactor: hook native Parquet writes into Spark's WriteFilesExec seam

docs/source/user-guide/latest/installation.md

@@ -142,8 +142,8 @@ Comet will log output similar to:
 ```shell
 INFO core/src/lib.rs: Comet native library version $COMET_VERSION initialized
 WARN CometExecRule: Comet cannot execute some parts of this plan natively (set spark.comet.explain.fallback.enabled=false to disable this logging):
-  Execute InsertIntoHadoopFsRelationCommand [COMET: Native support for operator DataWritingCommandExec is disabled. Set spark.comet.parquet.write.enabled=true to enable it.]
-+- WriteFiles
+  Execute InsertIntoHadoopFsRelationCommand
++- WriteFiles [COMET: Native support for operator WriteFilesExec is disabled. Set spark.comet.parquet.write.enabled=true to enable it.]
    +-  LocalTableScan [COMET: Native support for operator LocalTableScanExec is disabled. Set spark.comet.exec.localTableScan.enabled=true to enable it.]
 ```
 

docs/source/user-guide/latest/operators.md

@@ -116,9 +116,9 @@ omitted from the tables below and may be reconsidered based on demand:
 
 ## Writes
 
-| Operator                 | Status | Notes                                                             |
-| ------------------------ | ------ | ----------------------------------------------------------------- |
-| `DataWritingCommandExec` | ⚠️     | Experimental native Parquet writes, disabled by default (opt-in). |
+| Operator         | Status | Notes                                                                                                                             |
+| ---------------- | ------ | --------------------------------------------------------------------------------------------------------------------------------- |
+| `WriteFilesExec` | ⚠️     | Experimental native Parquet writes, disabled by default (opt-in). Requires Spark 4.0+. Non-partitioned, non-bucketed writes only. |
 
 ## Python and UDF
 

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

@@ -218,18 +218,11 @@ impl ParquetWriter {
 pub struct ParquetWriterExec {
     /// Input execution plan
     input: Arc<dyn ExecutionPlan>,
-    /// Output file path (final destination)
+    /// Path of the file to write. Chosen by the JVM commit protocol (see CometWriteFilesExec) and
+    /// used verbatim - this operator never derives file names of its own.
     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>,
     /// 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 +235,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>,
         compression: ParquetCompression,
-        partition_id: i32,
         column_names: Vec<String>,
         object_store_options: HashMap<String, String>,
     ) -> Result<Self> {
@@ -267,11 +255,7 @@ impl ParquetWriterExec {
         Ok(ParquetWriterExec {
             input,
             output_path,
-            work_dir,
-            job_id,
-            task_attempt_id,
             compression,
-            partition_id,
             column_names,
             object_store_options,
             metrics: ExecutionPlanMetricsSet::new(),
@@ -447,11 +431,7 @@ 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(),
                 self.object_store_options.clone(),
             )?)),
@@ -476,8 +456,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();
 
@@ -492,16 +470,8 @@ 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)
-        };
+        // The JVM commit protocol already decided where this task writes.
+        let part_file = self.output_path.clone();
 
         // Configure writer properties
         let props = WriterProperties::builder()
@@ -851,18 +821,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.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
             ParquetCompression::None,
-            0, // partition_id
             column_names,
             HashMap::new(), // object_store_options
         )?;

native/core/src/execution/planner.rs

@@ -1837,15 +1837,7 @@ impl PhysicalPlanner {
                 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,
                     codec,
-                    self.partition,
                     writer.column_names.clone(),
                     object_store_options,
                 )?);

native/proto/src/proto/operator.proto

@@ -458,16 +458,17 @@ message ShuffleWriter {
 }
 
 message ParquetWriter {
+  // Fully-qualified path of the Parquet file that this task must write, set per task by
+  // CometWriteFilesExec 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.
   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;
+  // Was work_dir / job_id / task_attempt_id, used when the native writer derived its own file
+  // names from the task context. File naming now comes from the commit protocol instead.
+  reserved 5, 6, 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.

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -870,8 +870,13 @@ object CometConf extends ShimCometConf {
     .booleanConf
     .createWithEnvVarOrDefault("ENABLE_COMET_STRICT_TESTING", false)
 
-  val COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT: ConfigEntry[Boolean] =
-    createOperatorIncompatConfig("DataWritingCommandExec")
+  val COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT: ConfigEntry[Boolean] =
+    // Native writes used to replace the whole DataWritingCommandExec; they now replace only
+    // WriteFilesExec, so the opt-in moved with it. Keep the old key working for anyone who had
+    // already enabled the experimental writer.
+    createOperatorIncompatConfig(
+      "WriteFilesExec",
+      Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec")))
 
   /** Create a config to enable a specific operator */
   private def createExecEnabledConfig(
@@ -896,15 +901,18 @@ object CometConf extends ShimCometConf {
   private def configKeyToEnvVar(configKey: String): String =
     configKey.toUpperCase(Locale.ROOT).replace('.', '_')
 
-  private def createOperatorIncompatConfig(name: String): ConfigEntry[Boolean] = {
+  private def createOperatorIncompatConfig(
+      name: String,
+      alternative: Option[String] = None): ConfigEntry[Boolean] = {
     val configKey = getOperatorAllowIncompatConfigKey(name)
     val envVar = configKeyToEnvVar(configKey)
-    conf(configKey)
+    val builder = conf(configKey)
       .category(CATEGORY_EXEC)
       .doc(s"Whether to allow incompatibility for operator: $name. " +
         s"False by default. Can be overridden with $envVar env variable")
-      .booleanConf
-      .createWithEnvVarOrDefault(envVar, false)
+    // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself.
+    alternative.foreach(builder.withAlternative(_))
+    builder.booleanConf.createWithEnvVarOrDefault(envVar, false)
   }
 
   def isExprEnabled(name: String, conf: SQLConf = SQLConf.get): Boolean = {

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -35,7 +35,7 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec}
-import org.apache.spark.sql.execution.datasources.WriteFilesExec
+import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec}
 import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat
 import org.apache.spark.sql.execution.datasources.json.JsonFileFormat
 import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
@@ -101,6 +101,14 @@ object CometExecRule {
 
   val allExecs: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] = nativeExecs ++ sinks
 
+  /**
+   * Output path of the write that a `WriteFilesExec` belongs to, copied from the enclosing
+   * `InsertIntoHadoopFsRelationCommand`. `WriteFilesExec` itself has no output path, so this is
+   * how `CometWriteFiles` learns the target filesystem - and, by its absence, that a write comes
+   * from some other V1 write command.
+   */
+  val WRITE_OUTPUT_PATH: TreeNodeTag[String] = TreeNodeTag[String]("comet.writeOutputPath")
+
   /**
    * Tag set on a `ShuffleExchangeExec` that should be left as a plain Spark shuffle rather than
    * wrapped in `CometShuffleExchangeExec`. See `tagRedundantColumnarShuffle`.
@@ -287,16 +295,11 @@ case class CometExecRule(session: SparkSession)
       case op if shouldApplySparkToColumnar(conf, op) =>
         convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
 
-      // AQE reoptimization looks for `DataWritingCommandExec` or `WriteFilesExec`
-      // if there is none it would reinsert write nodes, and since Comet remap those nodes
-      // to Comet counterparties the write nodes are twice to the plan.
-      // Checking if AQE inserted another write Command on top of existing write command
-      case _ @DataWritingCommandExec(_, w: WriteFilesExec)
-          if w.child.isInstanceOf[CometNativeWriteExec] =>
-        w.child
-
-      case op: DataWritingCommandExec =>
-        convertToComet(op, CometDataWritingCommand).getOrElse(op)
+      // Replace only the per-task write with Comet, leaving DataWritingCommandExec (and therefore
+      // Spark's commit protocol, stats trackers and SaveMode handling) in place. CometWriteFiles
+      // declines the write on Spark 3.x, where the node cannot be replaced at all.
+      case w: WriteFilesExec =>
+        convertToComet(w, CometWriteFiles).getOrElse(w)
 
       // For AQE broadcast stage on a Comet broadcast exchange
       case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) =>
@@ -372,14 +375,14 @@ case class CometExecRule(session: SparkSession)
         op match {
           case _: CometPlan | _: AQEShuffleReadExec | _: BroadcastExchangeExec |
               _: BroadcastQueryStageExec | _: AdaptiveSparkPlanExec | _: ExecutedCommandExec |
-              _: V2CommandExec | _: WriteFilesExec =>
+              _: V2CommandExec | _: WriteFilesExec | _: DataWritingCommandExec =>
             // Some execs should never be replaced. We include
             // these cases specially here so we do not add a misleading 'info' message.
-            // WriteFilesExec is always wrapped by DataWritingCommandExec (via Spark's V1Writes
-            // rule); the parent case converts the whole write to CometNativeWriteExec and
-            // unwraps WriteFilesExec inside convertToComet. Tagging WriteFilesExec here would
-            // produce a spurious "WriteFilesExec is not supported" fallback reason (and a warning
-            // when COMET_EXPLAIN_FALLBACK_LOG_ENABLED=true) even when the write is fully native.
+            // DataWritingCommandExec is deliberately left in the plan even for a fully native
+            // write - Comet replaces only its WriteFilesExec child (see CometWriteFiles) - so
+            // tagging it would report an accelerated write as a fallback. WriteFilesExec that
+            // reaches this point was already offered to CometWriteFiles by the case above, so it
+            // has a fallback reason already and must not be tagged again.
             op
           case _ =>
             // The operator was not converted to a Comet plan and no serde handler claimed it, so
@@ -401,6 +404,17 @@ case class CometExecRule(session: SparkSession)
         }
     }
 
+    // `WriteFilesExec` does not carry the write's output path, but CometWriteFiles needs it to
+    // decide whether the target filesystem is supported. Record it from the enclosing command
+    // before the bottom-up walk reaches the write node. The absence of the tag also tells
+    // CometWriteFiles that the write is not an InsertIntoHadoopFsRelationCommand and must be
+    // declined.
+    plan.foreach {
+      case DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, w: WriteFilesExec) =>
+        w.setTagValue(CometExecRule.WRITE_OUTPUT_PATH, cmd.outputPath.toString)
+      case _ =>
+    }
+
     plan.transformUp { case op =>
       val converted = convertNode(op)
       // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions
@@ -681,14 +695,14 @@ case class CometExecRule(session: SparkSession)
             firstNativeOp = true
           }
 
-          // CometNativeWriteExec is special: it has two separate plans:
+          // CometWriteFilesExec is special: it has two separate plans:
           // 1. A protobuf plan (nativeOp) describing the write operation
           // 2. A Spark plan (child) that produces the data to write
           // The serializedPlanOpt is a def that always returns Some(...) by serializing
           // nativeOp on-demand, so it doesn't need convertBlock(). However, its child
           // (e.g., CometNativeScanExec) may need its own serialization. Reset the flag
           // so children can start their own native execution blocks.
-          if (op.isInstanceOf[CometNativeWriteExec]) {
+          if (op.isInstanceOf[CometWriteFilesExec]) {
             firstNativeOp = true
           }
 
@@ -725,12 +739,7 @@ case class CometExecRule(session: SparkSession)
       // children are CometNativeExec. This prevents runtime failures when the native operator
       // expects Arrow arrays but receives non-Arrow data (e.g., OnHeapColumnVector).
       if (serde.requiresNativeChildren && op.children.nonEmpty) {
-        // Get the actual data-producing children (unwrap WriteFilesExec if present)
-        val dataProducingChildren = op.children.flatMap {
-          case writeFiles: WriteFilesExec => Seq(writeFiles.child)
-          case other => Seq(other)
-        }
-        if (!dataProducingChildren.forall(_.isInstanceOf[CometNativeExec])) {
+        if (!op.children.forall(_.isInstanceOf[CometNativeExec])) {
           withFallbackReason(
             op,
             "Cannot perform native operation because input is not in Arrow format")

spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala

@@ -22,7 +22,7 @@ package org.apache.comet.rules
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.sideBySide
-import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec}
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch}
 import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan}
@@ -87,10 +87,6 @@ case class EliminateRedundantTransitions(session: SparkSession)
           // and CometSparkToColumnarExec
           sparkToColumnar.child
         }
-      // Remove unnecessary transition for native writes
-      // Write should be final operation in the plan
-      case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) =>
-        nativeWrite
       case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) =>
         val op = createColumnarToRowExec(child)
         if (c.logicalLink.isEmpty) {