feat: support the S3A magic committer in native Parquet writes

.github/workflows/pr_build_linux.yml

@@ -497,6 +497,7 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.parquet.CometParquetWriterRetrySuite
               org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite
               org.apache.spark.sql.comet.CometNativeWriteSuite
               org.apache.comet.parquet.ParquetReadV1Suite

.github/workflows/pr_build_macos.yml

@@ -163,6 +163,7 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.parquet.CometParquetWriterRetrySuite
               org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite
               org.apache.spark.sql.comet.CometNativeWriteSuite
               org.apache.comet.parquet.ParquetReadV1Suite

dev/ci/check-suites.py

@@ -35,6 +35,7 @@ def file_to_class_name(path: Path) -> str | None:
     ignore_list = [
         "org.apache.comet.parquet.ParquetReadSuite", # abstract
         "org.apache.comet.parquet.ParquetReadFromS3Suite", # manual test suite
+        "org.apache.comet.parquet.ParquetWriteWithMagicCommitterSuite", # manual test suite
         "org.apache.comet.csv.CsvReadFromS3Suite", # manual test suite
         "org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite", # manual test suite (loads libhdfs, see #5023)
         "org.apache.comet.IcebergReadFromS3Suite", # manual test suite

docs/source/contributor-guide/memory_management.md

@@ -414,7 +414,11 @@ the executor.
   in-progress row group (`ArrowWriter::memory_size`). The estimate counts encoded pages, encoder
   buffers, dictionaries and Bloom filters. The writer keeps parquet-rs's default row-group limit of
   1Mi rows and sets no byte limit, so a wide schema can hold a large row group
-  ([#5304](https://github.com/apache/datafusion-comet/issues/5304)).
+  ([#5304](https://github.com/apache/datafusion-comet/issues/5304)). A file that is neither local
+  nor on HDFS, such as one the S3A magic committer hands out, is written through a Hadoop output
+  stream, and the reservation then also includes the 1 MiB buffer that batches the writes to it.
+  The stream's own buffers, such as S3A's upload buffers (`fs.s3a.fast.upload.buffer`), are not
+  Comet's and sit outside its pool, as they do for Spark's writer.
 - The Iceberg writer keeps one file open per partition in a fanout write, so its reservation grows
   with the number of partitions a task writes. iceberg-rust's `ParquetWriter` does not expose
   parquet-rs's `memory_size`, only every byte the file has written, its in-progress row group's

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

@@ -135,8 +135,8 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx
 
 | 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.                                                                                                                                                                                                                                                       |
+| `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. See [Parquet](parquet.md#writing).                                                                                                                                                                                                       |
+| `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. See [Parquet](parquet.md#writing).                                                                                                                                                                                                                    |
 | Iceberg writes: `AppendDataExec`, `OverwriteByExpressionExec`, `OverwritePartitionsDynamicExec`, `ReplaceDataExec` | ⚠️     | Disabled by default (opt-in). `spark.comet.write.iceberg.enabled=true` plans the write as `IcebergWrite` and `IcebergCommit` and writes eligible data files natively (`CometIcebergWrite`). Other writes use Iceberg's Java writer. Covers `INSERT INTO`, `INSERT OVERWRITE`, and copy-on-write `DELETE` / `UPDATE` / `MERGE`. Merge-on-read writes (`WriteDeltaExec`) fall back. See [Iceberg Writes](iceberg.md#writes). |
 | `MergeRowsExec`                                                                                                    | ⚠️     | Spark 3.5+. Experimental, disabled by default (`spark.comet.exec.mergeRows.enabled`). On Spark 4.1+, stock V2 writers retain Spark MergeRowsExec; Comet's split Iceberg path can run it natively while preserving MergeSummary. See [MERGE INTO](compatibility/operators.md#merge-into-mergerowsexec).                                                                                                                     |
 

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

@@ -170,6 +170,14 @@ partitioned and bucketed writes fall back to Spark, as do writes with
 behavior are not reproduced, as listed under [Known Issues](#known-issues). Iceberg writes use a
 different native writer, described on the [Iceberg](iceberg.md#writes) page.
 
+The native writer supports local and HDFS destinations, and S3 writes that the S3A magic committer
+commits. Spark 4.1 configures that committer by default when spark-hadoop-cloud is on the
+classpath, and earlier versions for a bucket that sets
+`fs.s3a.bucket.<bucket>.committer.magic.enabled`. S3A creates each of those files, through the same
+call Spark's own writer makes, so the magic committer and S3A's write settings, such as server-side
+encryption, ACLs, storage class and object headers, apply exactly as they do for Spark's writer.
+Other S3 writes fall back to Spark.
+
 ## Tuning
 
 ### Filter Pushdown / Late Materialization

native/core/src/execution/jni_api.rs

@@ -25,9 +25,11 @@
 //! `createPlan`, `setShufflePartitionPusher`, `executePlan` and `releasePlan` also depend on
 //! upcalls into the JVM (input iterators, the task memory manager, metrics, UDFs and scalar
 //! subqueries). Their core functions (`create_execution_context`,
-//! `register_shuffle_partition_pusher`, `execute_plan` and `release_plan`) take the JVM objects
-//! as JNI global references held by the `ExecutionContext` (`PlanJvmRefs`), and the steps that
-//! need an `Env` as callbacks: the memory pool factory, the pusher factory, and `PublishMetrics`.
+//! `register_shuffle_partition_pusher`, `execute_plan` and `release_plan`) take the JVM objects as
+//! JNI global references held by the `ExecutionContext` (`PlanJvmRefs`), and the steps that need
+//! an `Env` as callbacks: the memory pool factory, the pusher factory, and `PublishMetrics`.
+//! `setOutputStream` registers a native write's output stream through the same
+//! `register_before_execution` check as the pusher.
 
 use super::{serde, utils::SparkArrowConvert};
 use crate::{
@@ -36,7 +38,10 @@ use crate::{
         metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype,
         shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort,
     },
-    jvm_bridge::{JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher},
+    jvm_bridge::{
+        JVMClasses, JavaOutputStreamSink, JavaShufflePartitionPusher, OutputStreamSink,
+        ShufflePartitionPusher,
+    },
 };
 use std::collections::HashSet;
 
@@ -583,6 +588,10 @@ struct ExecutionContext {
     /// Task-owned remote shuffle callback, registered before native planning starts.
     /// The callback owns a JNI global reference and can safely run on Tokio workers.
     pub shuffle_partition_pusher: Option<Arc<dyn ShufflePartitionPusher>>,
+    /// The output stream of the file this plan's Parquet writer writes, registered before native
+    /// planning starts when the JVM created that file through Hadoop's FileSystem API. Like the
+    /// shuffle callback, it owns a JNI global reference and can run on any native thread.
+    pub output_stream: Option<Arc<dyn OutputStreamSink>>,
     /// Removes this context's tracing memory-pool entry on every exit path.
     memory_pool_registration: Option<ThreadMemoryPoolRegistration>,
 }
@@ -891,6 +900,7 @@ fn create_execution_context(
         task_context: jvm.task_context,
         class_loader: jvm.class_loader,
         shuffle_partition_pusher: None,
+        output_stream: None,
         memory_pool_registration,
     }))
 }
@@ -929,23 +939,72 @@ fn register_shuffle_partition_pusher(
     exec_context: &mut ExecutionContext,
     make_pusher: impl FnOnce() -> CometResult<Arc<dyn ShufflePartitionPusher>>,
 ) -> CometResult<()> {
-    if exec_context.root_op.is_some() {
-        return Err(CometError::Internal(
-            "Remote shuffle callback cannot be registered after native execution starts"
-                .to_string(),
-        ));
+    register_before_execution(
+        exec_context.root_op.is_some(),
+        &mut exec_context.shuffle_partition_pusher,
+        "Remote shuffle callback",
+        make_pusher,
+    )
+}
+
+/// Registers one of a plan's task-owned JVM callbacks in `slot`, built by `make` only once the
+/// plan is known to accept it: before execution has `started`, and once. `what` names the callback
+/// in the errors.
+fn register_before_execution<T: ?Sized>(
+    started: bool,
+    slot: &mut Option<Arc<T>>,
+    what: &str,
+    make: impl FnOnce() -> CometResult<Arc<T>>,
+) -> CometResult<()> {
+    if started {
+        return Err(CometError::Internal(format!(
+            "{what} cannot be registered after native execution starts"
+        )));
     }
 
-    if exec_context.shuffle_partition_pusher.is_some() {
-        return Err(CometError::Internal(
-            "Remote shuffle callback has already been registered for this task".to_string(),
-        ));
+    if slot.is_some() {
+        return Err(CometError::Internal(format!(
+            "{what} has already been registered for this task"
+        )));
     }
 
-    exec_context.shuffle_partition_pusher = Some(make_pusher()?);
+    *slot = Some(make()?);
     Ok(())
 }
 
+/// Binds the output stream of the file the plan's Parquet writer writes, for a file the JVM
+/// created itself, before native execution is initialized.
+///
+/// Separate from `createPlan` for the same reason as `setShufflePartitionPusher`: only native
+/// writes the JVM creates the file for need it, so every other caller keeps the existing ABI.
+#[no_mangle]
+pub extern "system" fn Java_org_apache_comet_Native_setOutputStream(
+    e: EnvUnowned,
+    _class: JClass,
+    exec_context: jlong,
+    output: JObject,
+) {
+    try_unwrap_or_throw(&e, |env| {
+        if exec_context == 0 {
+            return Err(CometError::NullPointer(
+                "A native write's output stream requires a valid native execution plan".to_string(),
+            ));
+        }
+
+        let exec_context = get_execution_context(exec_context);
+        register_before_execution(
+            exec_context.root_op.is_some(),
+            &mut exec_context.output_stream,
+            "A native write's output stream",
+            || {
+                let sink: Arc<dyn OutputStreamSink> =
+                    Arc::new(JavaOutputStreamSink::try_new(env, &output)?);
+                Ok(sink)
+            },
+        )
+    })
+}
+
 /// Skipping is opt-in (`spark.comet.exec.aggregate.skipPartial.enabled`): once DataFusion's probe
 /// decides to skip, it never aggregates again, so a task whose keys repeat after a mostly distinct
 /// start shuffles every later row (#6466).
@@ -1467,7 +1526,8 @@ fn execute_plan(
             .with_sql_text_pool(&exec_context.spark_plan)
             .with_task_context(exec_context.task_context.clone())
             .with_class_loader(exec_context.class_loader.clone())
-            .with_shuffle_partition_pusher(exec_context.shuffle_partition_pusher.clone());
+            .with_shuffle_partition_pusher(exec_context.shuffle_partition_pusher.clone())
+            .with_output_stream(exec_context.output_stream.clone());
         let (scans, shuffle_scans, root_op) = planner.create_plan(
             &exec_context.spark_plan,
             &mut exec_context.input_sources.clone(),