feat: hook native Parquet writes into Spark's WriteFilesExec seam on Spark 4.0+

benchmarks/pyspark/run_all_benchmarks.sh

@@ -78,6 +78,7 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.memory.offHeap.enabled=true \
   --conf spark.memory.offHeap.size=16g \
   --conf spark.comet.enabled=true \
+  --conf spark.comet.operator.WriteFilesExec.allowIncompatible=true \
   --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \
   --conf spark.comet.parquet.write.enabled=true \
   --conf spark.comet.explain.fallback.log.enabled=true \
@@ -106,6 +107,7 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.memory.offHeap.enabled=true \
   --conf spark.memory.offHeap.size=16g \
   --conf spark.comet.enabled=true \
+  --conf spark.comet.operator.WriteFilesExec.allowIncompatible=true \
   --conf spark.comet.operator.DataWritingCommandExec.allowIncompatible=true \
   --conf spark.comet.parquet.write.enabled=true \
   --conf spark.comet.explain.fallback.log.enabled=true \

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

@@ -181,16 +181,20 @@ Create a test Parquet source
 scala> (0 until 10).toDF("a").write.mode("overwrite").parquet("/tmp/test")
 ```
 
-Comet will log output similar to:
+Comet will log output similar to this on Spark 4.0 and later:
 
 ```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.]
 ```
 
+On Spark 3.4 and 3.5 the native writer replaces the whole write command rather than just the
+per-task write, so the same message appears on `Execute InsertIntoHadoopFsRelationCommand` and
+names `DataWritingCommandExec`.
+
 Query the data from the test source and check:
 
 - INFO message shows the native Comet library has been initialized.

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

@@ -117,9 +117,10 @@ 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`         | ⚠️     | Spark 4.0+. Experimental native Parquet writes, disabled by default (opt-in). Non-partitioned, non-bucketed writes only, and not when `spark.sql.files.maxRecordsPerFile` is set. |
+| `DataWritingCommandExec` | ⚠️     | Spark 3.4/3.5 only. Experimental native Parquet writes, disabled by default (opt-in). Replaced by `WriteFilesExec` on Spark 4.0+ and removed with Spark 3.x support.              |
 
 ## Python and UDF
 

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

@@ -220,10 +220,14 @@ impl ParquetWriter {
 pub struct ParquetWriterExec {
     /// Input execution plan
     input: Arc<dyn ExecutionPlan>,
-    /// Output file path (final destination)
+    /// 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.
     output_path: String,
-    /// Working directory for temporary files (used by FileCommitProtocol)
-    work_dir: 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
@@ -250,7 +254,7 @@ impl ParquetWriterExec {
     pub fn try_new(
         input: Arc<dyn ExecutionPlan>,
         output_path: String,
-        work_dir: String,
+        work_dir: Option<String>,
         job_id: Option<String>,
         task_attempt_id: Option<i32>,
         compression: ParquetCompression,
@@ -510,15 +514,18 @@ impl ExecutionPlan for ParquetWriterExec {
             Arc::new(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)
+        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),
+            },
         };
 
         // Configure writer properties
@@ -647,6 +654,56 @@ 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
+    /// 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<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+        )?;
+
+        let memory_source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None)?;
+        let input = Arc::new(DataSourceExec::new(Arc::new(memory_source)));
+        let temp_dir = tempfile::tempdir()?;
+        // A name the writer could never have derived itself, matching Spark's convention.
+        let file_name = "part-00007-11111111-2222-3333-4444-555555555555-c000.parquet";
+        let output_path = format!("file://{}/{}", temp_dir.path().display(), file_name);
+
+        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,
+            vec!["id".to_string()],
+            None,
+            HashMap::new(),
+        )?;
+
+        let mut stream = writer.execute(0, SessionContext::new().task_ctx())?;
+        while stream.try_next().await?.is_some() {}
+
+        let written = temp_dir.path().join(file_name);
+        assert!(
+            written.exists(),
+            "expected the writer to use the given path verbatim, found: {:?}",
+            std::fs::read_dir(temp_dir.path())?
+                .filter_map(|e| e.ok().map(|e| e.file_name()))
+                .collect::<Vec<_>>()
+        );
+
+        let reader = SerializedFileReader::new(File::open(written)?)?;
+        assert_eq!(reader.metadata().file_metadata().num_rows(), 3);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_parquet_writer_preserves_catalyst_schema_in_footer() -> Result<()> {
         let values = ListArray::from_iter_primitive::<Int32Type, _, _>([
@@ -689,7 +746,7 @@ mod tests {
         let writer = ParquetWriterExec::try_new(
             input,
             work_dir.clone(),
-            work_dir,
+            Some(work_dir),
             None,
             None,
             ParquetCompression::None,
@@ -758,7 +815,7 @@ mod tests {
         let writer = ParquetWriterExec::try_new(
             input,
             work_dir.clone(),
-            work_dir,
+            Some(work_dir),
             None,
             None,
             ParquetCompression::None,
@@ -1018,7 +1075,7 @@ mod tests {
         let parquet_writer = ParquetWriterExec::try_new(
             memory_exec,
             output_path,
-            work_dir,
+            Some(work_dir),
             None,      // job_id
             Some(123), // task_attempt_id
             ParquetCompression::None,

native/core/src/execution/planner.rs

@@ -1991,11 +1991,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.work_dir.clone(),
                     writer.job_id.clone(),
                     writer.task_attempt_id,
                     codec,

native/core/src/parquet/parquet_support.rs

@@ -1277,6 +1277,53 @@ mod tests {
         );
     }
 
+    /// Pins the set of characters the `url` crate rewrites inside a path.
+    ///
+    /// `create_hdfs_object_store` hands `url.path()` to `object_store::path::Path::parse`, so a
+    /// character that the parser escapes - or, for `?` and `#`, treats as a delimiter and
+    /// truncates at - makes the native writer create a different file from the one Spark's commit
+    /// protocol chose and later commits. Job commit still succeeds, so the data simply ends up
+    /// somewhere nobody looks.
+    ///
+    /// `NativeWriteUtils.nativeUrlEscapedAscii` on the JVM side keeps a copy of this set and
+    /// declines such destinations during planning. Nothing else would notice if a `url` upgrade
+    /// moved a character between the two groups and quietly reopened that hole, so assert the set
+    /// here rather than in the Scala test, which can only check the copy against itself.
+    #[test]
+    fn url_path_rewritten_characters() {
+        use url::Url;
+
+        let rewritten: String = (' '..='~')
+            .filter(|c| {
+                let url = Url::parse(&format!("hdfs://ns/pre{c}post/output")).unwrap();
+                url.path() != format!("/pre{c}post/output")
+            })
+            .collect();
+
+        assert_eq!(
+            rewritten, " \"#<>?`{}",
+            "the ASCII characters `url` rewrites inside a path changed; update \
+             NativeWriteUtils.nativeUrlEscapedAscii and its tests to match"
+        );
+
+        // Every non-ASCII byte is percent-encoded regardless of the encode set. This is the half
+        // a `java.net.URI` raw/decoded comparison cannot see, because Java leaves non-ASCII path
+        // characters alone.
+        for name in [
+            "caf\u{e9}",                // precomposed e-acute
+            "cafe\u{301}",              // "e" plus a combining acute accent
+            "\u{65e5}\u{672c}\u{8a9e}", // CJK
+            "\u{1f642}",                // astral plane
+        ] {
+            let url = Url::parse(&format!("hdfs://ns/{name}/output")).unwrap();
+            assert_ne!(
+                url.path(),
+                format!("/{name}/output"),
+                "expected `url` to percent-encode the non-ASCII name {name}"
+            );
+        }
+    }
+
     #[cfg(not(feature = "hdfs-opendal"))]
     #[test]
     fn test_prepare_object_store() {

native/proto/src/proto/operator.proto

@@ -897,11 +897,22 @@ 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.
   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
+  // 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;

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

@@ -1024,9 +1024,29 @@ object CometConf extends ShimCometConf {
     .booleanConf
     .createWithEnvVarOrDefault("ENABLE_COMET_STRICT_TESTING", false)
 
+  /**
+   * Deprecated alternatives for `spark.comet.operator.<name>.allowIncompatible`, keyed by
+   * operator name. `CometExecRule` resolves these configs by operator name rather than through
+   * the registered `ConfigEntry`, so `isOperatorAllowIncompat` has to consult the alternatives
+   * itself - otherwise an old key would read `true` from the entry while the planner saw `false`.
+   */
+  private val operatorIncompatAlternatives =
+    scala.collection.mutable.Map.empty[String, String]
+
+  /** Spark 3.x only: native writes replace the whole `DataWritingCommandExec` there. */
   val COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT: ConfigEntry[Boolean] =
     createOperatorIncompatConfig("DataWritingCommandExec")
 
+  /**
+   * Spark 4.0+ only: native writes replace just the `WriteFilesExec` child, so the opt-in moved
+   * with the operator. The old `DataWritingCommandExec` key keeps working for anyone who had
+   * already enabled the experimental writer.
+   */
+  val COMET_OPERATOR_WRITE_FILES_ALLOW_INCOMPAT: ConfigEntry[Boolean] =
+    createOperatorIncompatConfig(
+      "WriteFilesExec",
+      Some(getOperatorAllowIncompatConfigKey("DataWritingCommandExec")))
+
   /** Create a config to enable a specific operator */
   private def createExecEnabledConfig(
       exec: String,
@@ -1050,15 +1070,21 @@ 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)
+    alternative.foreach { alt =>
+      operatorIncompatAlternatives.put(name, alt)
+      // ConfigBuilder mutates in place, so the result of withAlternative is `builder` itself.
+      builder.withAlternative(alt)
+    }
+    builder.booleanConf.createWithEnvVarOrDefault(envVar, false)
   }
 
   def isExprEnabled(name: String, conf: SQLConf = SQLConf.get): Boolean = {
@@ -1082,7 +1108,11 @@ object CometConf extends ShimCometConf {
   }
 
   def isOperatorAllowIncompat(name: String, conf: SQLConf = SQLConf.get): Boolean = {
-    getBooleanConf(getOperatorAllowIncompatConfigKey(name), defaultValue = false, conf)
+    val value = CometConfDeprecations.readWithAlternatives(
+      conf,
+      getOperatorAllowIncompatConfigKey(name),
+      operatorIncompatAlternatives.get(name).toSeq)
+    value != null && value.toLowerCase(Locale.ROOT) == "true"
   }
 
   def getOperatorAllowIncompatConfigKey(name: String): String = {