From c8996d5107641f1d8e337951fa11fa00693fe92d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 13 Oct 2026 12:58:15 -0600 Subject: [PATCH 1/7] feat: support the S3A magic committer in native Parquet writes On Spark 4.0+, a task file that is neither local nor on HDFS is created by the JVM through Hadoop's FileSystem API, as Spark's Parquet writer creates it, and the native writer streams the file's bytes into that stream over JNI. S3A therefore creates the magic committer's files itself and records their pending uploads. A failed write aborts the stream. Planning admits an S3 destination when the S3A magic committer commits it, through PathOutputCommitProtocol or BindingParquetOutputCommitter. Other S3 writes still fall back to Spark. --- dev/ci/check-suites.py | 1 + docs/source/user-guide/latest/operators.md | 7 + native/core/src/execution/jni_api.rs | 120 ++++++- .../src/execution/operators/parquet_writer.rs | 328 ++++++++++++++++-- native/core/src/execution/planner.rs | 21 +- native/core/src/execution/planner/write.rs | 73 +++- native/jni-bridge/src/lib.rs | 2 + native/jni-bridge/src/output_stream_sink.rs | 141 ++++++++ native/proto/src/proto/operator.proto | 8 +- pom.xml | 38 ++ spark/pom.xml | 4 + .../comet/parquet/HadoopFileOutput.java | 126 +++++++ .../org/apache/comet/CometExecIterator.scala | 7 +- .../main/scala/org/apache/comet/Native.scala | 16 +- .../serde/operator/CometWriteFiles.scala | 16 +- .../serde/operator/NativeWriteUtils.scala | 90 +++++ .../spark/sql/comet/CometWriteFilesExec.scala | 85 +++-- .../comet/hadoop/fs/RecordingFileSystem.java | 190 ++++++++++ .../parquet/CometParquetWriterSuite.scala | 205 ++++++++++- .../ParquetWriteWithMagicCommitterSuite.scala | 249 +++++++++++++ 20 files changed, 1664 insertions(+), 63 deletions(-) create mode 100644 native/jni-bridge/src/output_stream_sink.rs create mode 100644 spark/src/main/java/org/apache/comet/parquet/HadoopFileOutput.java create mode 100644 spark/src/test/java/org/apache/comet/hadoop/fs/RecordingFileSystem.java create mode 100644 spark/src/test/scala/org/apache/comet/parquet/ParquetWriteWithMagicCommitterSuite.scala diff --git a/dev/ci/check-suites.py b/dev/ci/check-suites.py index 5b6e0f67e3b..1c30e3b828f 100644 --- a/dev/ci/check-suites.py +++ b/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.parquet.ParquetReadFromFakeHadoopFsSuite", # manual test suite (loads libhdfs, see #5023) "org.apache.comet.IcebergReadFromS3Suite", # manual test suite "org.apache.comet.cloud.s3.CometS3CredentialBridgeSuite", # manual test suite diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index c0054d2213b..8f3911f8af2 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -139,6 +139,13 @@ natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-m | Iceberg writes: `AppendDataExec`, `OverwriteByExpressionExec`, `OverwritePartitionsDynamicExec`, `ReplaceDataExec` | ⚠️ | Experimental, disabled by default. `spark.comet.write.iceberg.splitOperator.enabled=true` plans the write as `IcebergWrite` and `IcebergCommit`, and Iceberg's Java writer still writes the data files. `spark.comet.write.iceberg.enabled=true` then writes eligible data files natively (`CometIcebergWrite`). Covers `INSERT INTO`, `INSERT OVERWRITE`, and copy-on-write `DELETE` / `UPDATE` / `MERGE`. Merge-on-read writes (`WriteDeltaExec`) fall back. See [Iceberg Writes](iceberg-writes.md). | | `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). | +Native Parquet writes support local and HDFS destinations. On Spark 4.0 and later they also support +S3 writes that the S3A magic committer commits, which Spark 4.1 configures by default when +spark-hadoop-cloud is on the classpath. S3A creates each of those files, through the same call +Spark's own writer makes, and Comet's native writer only produces the Parquet bytes. The magic +committer and S3A's write settings, such as server-side encryption, ACLs, storage class and object +headers, therefore apply exactly as they do for Spark's writer. Other S3 writes fall back to Spark. + ## Python and UDF | Operator | Status | Notes | diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 262d8b8fd41..7974ae3f082 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -22,12 +22,13 @@ //! converts the result back. Errors and panics are raised as JVM exceptions by //! `try_unwrap_or_throw`, through the boundary error protocol in `errors::NativeError`. //! -//! `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`. +//! `createPlan`, `setShufflePartitionPusher`, `setOutputStream`, `executePlan` and `releasePlan` +//! also depend on upcalls into the JVM (input iterators, the task memory manager, metrics, UDFs, +//! scalar subqueries and a native write's output stream). Their core functions +//! (`create_execution_context`, `register_shuffle_partition_pusher`, `register_output_stream`, +//! `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 and output stream factories, and `PublishMetrics`. use super::{serde, utils::SparkArrowConvert}; use crate::{ @@ -36,7 +37,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 +587,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>, + /// 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>, /// Removes this context's tracing memory-pool entry on every exit path. memory_pool_registration: Option, } @@ -891,6 +899,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, })) } @@ -946,6 +955,56 @@ fn register_shuffle_partition_pusher( 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_output_stream(exec_context, || { + Ok(Arc::new(JavaOutputStreamSink::try_new(env, &output)?)) + }) + }) +} + +/// Registers the output stream of the file the plan's Parquet writer writes, built by +/// `make_sink` only once the plan is known to accept one: before execution starts, and once. Core +/// of `Native.setOutputStream`, whose sink calls back into the JVM. +fn register_output_stream( + exec_context: &mut ExecutionContext, + make_sink: impl FnOnce() -> CometResult>, +) -> CometResult<()> { + if exec_context.root_op.is_some() { + return Err(CometError::Internal( + "A native write's output stream cannot be registered after native execution starts" + .to_string(), + )); + } + + if exec_context.output_stream.is_some() { + return Err(CometError::Internal( + "A native write's output stream has already been registered for this plan".to_string(), + )); + } + + exec_context.output_stream = Some(make_sink()?); + Ok(()) +} + /// 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(), @@ -3789,4 +3849,48 @@ mod tests { ); release_plan(exec_context, &mut |_| Ok(())).unwrap(); } + + struct DiscardingSink; + + impl OutputStreamSink for DiscardingSink { + fn max_write_size(&self) -> usize { + 1 + } + + fn write(&self, _data: &[u8]) -> datafusion::common::Result<()> { + Ok(()) + } + } + + #[test] + fn register_output_stream_once_before_execution() { + let _guard = serial(); + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + register_output_stream(&mut exec_context, || Ok(Arc::new(DiscardingSink))).unwrap(); + assert!(exec_context.output_stream.is_some()); + + // The sink is built only once the plan is known to accept it. + let error = register_output_stream(&mut exec_context, || { + panic!("a rejected output stream must not be built") + }) + .unwrap_err(); + assert!( + error.to_string().contains("already been registered"), + "{error}" + ); + + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap(); + let error = register_output_stream(&mut exec_context, || { + panic!("a rejected output stream must not be built") + }) + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot be registered after native execution starts"), + "{error}" + ); + release_plan(exec_context, &mut |_| Ok(())).unwrap(); + } } diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 81fed5e5994..34347b8db1c 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -22,6 +22,7 @@ use std::{ fmt, fmt::{Debug, Formatter}, fs::File, + io::Write, sync::Arc, }; @@ -30,6 +31,7 @@ use opendal::Operator; #[cfg(feature = "hdfs-opendal")] use std::io::Cursor; +use crate::jvm_bridge::OutputStreamSink; use crate::parquet::parquet_support::is_hdfs_scheme; #[cfg(feature = "hdfs-opendal")] use crate::parquet::parquet_support::{create_hdfs_operator, prepare_object_store_with_configs}; @@ -99,6 +101,10 @@ enum ParquetWriter { Box, String, ), + /// Writer for a file the JVM created through Hadoop's FileSystem API, which only supplies the + /// file's bytes. Whatever the file system does to create and store the file, S3A's magic + /// committer support among it, is the JVM's. + OutputStream(ArrowWriter), } impl ParquetWriter { @@ -154,13 +160,15 @@ impl ParquetWriter { Ok(()) } + ParquetWriter::OutputStream(writer) => writer.write(batch), } } /// Memory the writer holds between batches: parquet-rs's estimate for the in-progress row /// group, which counts encoded pages, encoder buffers, dictionaries and any Bloom filters. /// The remote writer flushes a row group after every batch and stages it in a buffer that is - /// cleared after each upload but keeps its capacity, so that buffer counts too. + /// cleared after each upload but keeps its capacity, so that buffer counts too, as does the + /// buffer an output stream writer collects its next write in. fn memory_size(&self) -> usize { match self { ParquetWriter::LocalFile(writer) => writer.memory_size(), @@ -168,15 +176,20 @@ impl ParquetWriter { ParquetWriter::Remote(writer, ..) => { writer.memory_size() + writer.inner().get_ref().capacity() } + ParquetWriter::OutputStream(writer) => { + writer.memory_size() + writer.inner().buffer.capacity() + } } } - /// Close the writer and finalize the file - async fn close(self) -> std::result::Result<(), parquet::errors::ParquetError> { + /// Close the writer and finalize the file. Returns the file's size in bytes when the writer + /// counted it, which only an output stream writer does, since nothing else can tell the size + /// of a file the JVM created. Otherwise the caller has to ask the file system. + async fn close(self) -> std::result::Result, parquet::errors::ParquetError> { match self { ParquetWriter::LocalFile(writer) => { writer.close()?; - Ok(()) + Ok(None) } #[cfg(feature = "hdfs-opendal")] ParquetWriter::Remote( @@ -223,12 +236,79 @@ impl ParquetWriter { } } - Ok(()) + Ok(None) + } + ParquetWriter::OutputStream(mut writer) => { + // `finish` sends the footer and flushes the writer's last write, so the JVM can + // close the file once this returns. Unlike `close` it keeps the writer, which has + // counted every byte it wrote. + writer.finish()?; + Ok(Some(writer.bytes_written() as u64)) } } } } +/// Writes a Parquet file's bytes to the output stream of the file the JVM created. Bytes are +/// collected into writes of the stream's `max_write_size`, so a file costs one upcall per write +/// rather than one per page. +struct OutputStreamWriter { + sink: Arc, + /// Bytes not yet handed to the stream, at most `capacity` of them. + buffer: Vec, + capacity: usize, + /// Set once a write to the stream has failed. Nothing is sent after that, not even what + /// dropping the Parquet writer flushes, since the stream is in an unknown state and the JVM + /// aborts the file once the failure reaches it. + failed: bool, +} + +impl OutputStreamWriter { + fn new(sink: Arc) -> Self { + let capacity = sink.max_write_size().max(1); + Self { + sink, + buffer: Vec::with_capacity(capacity), + capacity, + failed: false, + } + } + + /// Hands the buffered bytes to the stream. + fn send(&mut self) -> std::io::Result<()> { + if self.failed { + return Err(std::io::Error::other( + "a previous write to the output stream failed", + )); + } + if !self.buffer.is_empty() { + if let Err(e) = self.sink.write(&self.buffer) { + self.failed = true; + return Err(std::io::Error::other(e)); + } + self.buffer.clear(); + } + Ok(()) + } +} + +impl Write for OutputStreamWriter { + fn write(&mut self, data: &[u8]) -> std::io::Result { + // A full buffer is sent before taking more rather than as soon as it fills, so that a + // failed send consumes none of `data`, as `Write` requires of an error. + if self.buffer.len() == self.capacity || self.failed { + self.send()?; + } + let taken = data.len().min(self.capacity - self.buffer.len()); + self.buffer.extend_from_slice(&data[..taken]); + Ok(taken) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.send() + } +} + /// Parquet writer operator that writes input batches to a Parquet file #[derive(Debug)] pub struct ParquetWriterExec { @@ -247,6 +327,9 @@ pub struct ParquetWriterExec { output_schema: Option, /// Object store configuration options object_store_options: HashMap, + /// The output stream of the file at `output_path`, when the JVM created that file. The writer + /// then writes the file's bytes there and never opens `output_path` itself. + output_stream: Option>, /// Metrics metrics: ExecutionPlanMetricsSet, /// Cache for plan properties @@ -283,11 +366,19 @@ impl ParquetWriterExec { column_names, output_schema, object_store_options, + output_stream: None, metrics: ExecutionPlanMetricsSet::new(), cache, }) } + /// Write the file's bytes to `output_stream`, the output stream of the file the JVM created + /// at `output_path`, rather than creating the file natively. + pub fn with_output_stream(mut self, output_stream: Arc) -> Self { + self.output_stream = Some(output_stream); + self + } + /// Create an Arrow writer based on the storage scheme /// /// # Arguments @@ -296,6 +387,7 @@ impl ParquetWriterExec { /// * `props` - Writer properties including compression /// * `runtime_env` - Runtime environment for object store registration /// * `object_store_options` - Configuration options for object store + /// * `output_stream` - The output stream of the file, when the JVM created it /// /// # Returns /// * `Ok(ParquetWriter)` - A writer appropriate for the storage scheme @@ -306,7 +398,22 @@ impl ParquetWriterExec { props: WriterProperties, _runtime_env: Arc, object_store_options: &HashMap, + output_stream: Option>, ) -> Result { + // The JVM created the file through its own FileSystem, so the path is not this writer's + // to interpret, whatever its scheme. + if let Some(output_stream) = output_stream { + let writer = + ArrowWriter::try_new(OutputStreamWriter::new(output_stream), schema, Some(props)) + .map_err(|e| { + DataFusionError::Execution(format!( + "Failed to create output stream writer: {}", + e + )) + })?; + return Ok(ParquetWriter::OutputStream(writer)); + } + // Parse URL and match on storage scheme directly let url = Url::parse(output_file_path).map_err(|e| { DataFusionError::Execution(format!("Failed to parse URL '{}': {}", output_file_path, e)) @@ -460,15 +567,21 @@ impl ExecutionPlan for ParquetWriterExec { children: Vec>, ) -> Result> { match children.len() { - 1 => Ok(Arc::new(ParquetWriterExec::try_new( - Arc::clone(&children[0]), - self.output_path.clone(), - self.compression.clone(), - self.partition_id, - self.column_names.clone(), - self.output_schema.clone(), - self.object_store_options.clone(), - )?)), + 1 => { + let writer = ParquetWriterExec::try_new( + Arc::clone(&children[0]), + self.output_path.clone(), + self.compression.clone(), + self.partition_id, + self.column_names.clone(), + self.output_schema.clone(), + self.object_store_options.clone(), + )?; + Ok(Arc::new(match &self.output_stream { + Some(output_stream) => writer.with_output_stream(Arc::clone(output_stream)), + None => writer, + })) + } _ => Err(DataFusionError::Internal( "ParquetWriterExec requires exactly one child".to_string(), )), @@ -525,6 +638,7 @@ impl ExecutionPlan for ParquetWriterExec { props, runtime_env, &object_store_options, + self.output_stream.clone(), )?; // Clone schema for use in async closure @@ -567,18 +681,23 @@ impl ExecutionPlan for ParquetWriterExec { reservation.try_resize(writer.memory_size())?; } - writer.close().await.map_err(|e| { + let written_size = writer.close().await.map_err(|e| { DataFusionError::Execution(format!("Failed to close writer: {}", e)) })?; - // Get file size - strip file:// prefix if present for local filesystem access - let local_path = part_file - .strip_prefix("file://") - .or_else(|| part_file.strip_prefix("file:")) - .unwrap_or(&part_file); - let file_size = std::fs::metadata(local_path) - .map(|m| m.len() as i64) - .unwrap_or(0); + let file_size = match written_size { + Some(size) => size as i64, + None => { + // Get file size - strip file:// prefix if present for local filesystem access + let local_path = part_file + .strip_prefix("file://") + .or_else(|| part_file.strip_prefix("file:")) + .unwrap_or(&part_file); + std::fs::metadata(local_path) + .map(|m| m.len() as i64) + .unwrap_or(0) + } + }; // Update metrics with write statistics files_written.add(1); @@ -917,6 +1036,168 @@ mod tests { Ok(()) } + /// Records what a writer hands the output stream of a file the JVM created, and fails the + /// write with index `failing_write` when there is one. + #[derive(Default)] + struct RecordingSink { + max_write_size: usize, + failing_write: Option, + writes: std::sync::Mutex>>, + } + + impl RecordingSink { + fn new(max_write_size: usize) -> Self { + Self { + max_write_size, + ..Default::default() + } + } + + fn writes(&self) -> Vec> { + self.writes.lock().unwrap().clone() + } + } + + impl OutputStreamSink for RecordingSink { + fn max_write_size(&self) -> usize { + self.max_write_size + } + + fn write(&self, data: &[u8]) -> Result<()> { + let mut writes = self.writes.lock().unwrap(); + assert!( + data.len() <= self.max_write_size, + "a write of {} bytes exceeds the stream's limit of {}", + data.len(), + self.max_write_size + ); + writes.push(data.to_vec()); + if self.failing_write == Some(writes.len() - 1) { + return Err(DataFusionError::Execution( + "injected output stream failure".to_string(), + )); + } + Ok(()) + } + } + + /// Writes `batches` with a writer for `output_path`, through `output_stream` when there is + /// one, and returns the writer for its metrics. + async fn write_batches( + batches: Vec, + output_path: String, + output_stream: Option>, + ) -> Result> { + let schema = batches[0].schema(); + let column_names = schema.fields().iter().map(|f| f.name().clone()).collect(); + let memory_source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let writer = ParquetWriterExec::try_new( + input, + output_path, + ParquetCompression::Snappy, + 0, + column_names, + None, + HashMap::new(), + )?; + let writer = Arc::new(match output_stream { + Some(output_stream) => writer.with_output_stream(output_stream), + None => writer, + }); + let mut stream = writer.execute(0, SessionContext::new().task_ctx())?; + while stream.try_next().await?.is_some() {} + Ok(writer) + } + + /// A file the JVM created gets exactly the bytes the writer would have written to a file of + /// its own, in writes no larger than the stream accepts, and the writer never opens the path + /// itself: here an S3A magic committer path, which only S3A can create correctly. + #[tokio::test] + async fn test_parquet_writer_writes_a_jvm_created_file_through_its_output_stream() -> Result<()> + { + let batches: Vec = (1..=5) + .map(create_test_record_batch) + .collect::>()?; + let temp_dir = tempfile::tempdir()?; + let local_file = temp_dir.path().join("part-00000.parquet"); + write_batches( + batches.clone(), + format!("file://{}", local_file.display()), + None, + ) + .await?; + let expected = std::fs::read(&local_file)?; + + let max_write_size = 4096; + let sink = Arc::new(RecordingSink::new(max_write_size)); + let magic_path = "s3a://bucket/table/__magic_job-0001/job-0001/tasks/\ + attempt_0001_m_000000_0/__base/part-00000.parquet"; + let writer = write_batches( + batches, + magic_path.to_string(), + Some(Arc::clone(&sink) as Arc), + ) + .await?; + + let writes = sink.writes(); + assert!( + writes.len() > 1, + "expected several writes, got {}", + writes.len() + ); + let (last, full) = writes.split_last().unwrap(); + assert!( + full.iter().all(|write| write.len() == max_write_size), + "every write but the last fills the stream's limit: {:?}", + writes.iter().map(Vec::len).collect::>() + ); + assert!(!last.is_empty() && last.len() <= max_write_size); + assert_eq!(writes.concat(), expected); + + let metrics = writer.metrics().unwrap(); + assert_eq!( + metrics.sum_by_name("bytes_written").unwrap().as_usize(), + expected.len() + ); + assert_eq!( + metrics.sum_by_name("rows_written").unwrap().as_usize(), + 5000 + ); + Ok(()) + } + + /// A failed write to the output stream fails the task with the stream's error, and nothing is + /// sent to the stream after it, not even what dropping the Parquet writer flushes. + #[tokio::test] + async fn test_parquet_writer_stops_writing_to_a_failed_output_stream() -> Result<()> { + let batches: Vec = (1..=5) + .map(create_test_record_batch) + .collect::>()?; + let sink = Arc::new(RecordingSink { + failing_write: Some(1), + ..RecordingSink::new(1024) + }); + let error = write_batches( + batches, + "s3a://bucket/table/part-00000.parquet".to_string(), + Some(Arc::clone(&sink) as Arc), + ) + .await + .expect_err("a failed write to the output stream must fail the task"); + + assert!( + error.to_string().contains("injected output stream failure"), + "{error}" + ); + assert_eq!( + sink.writes().len(), + 2, + "the writer kept writing to a stream that had failed" + ); + Ok(()) + } + #[tokio::test] #[cfg(feature = "hdfs-opendal")] #[ignore = "This test requires a running HDFS cluster"] @@ -1068,6 +1349,7 @@ mod tests { props, runtime_env, &HashMap::new(), + None, )?; // Write 5 batches in a loop diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a05de76019d..ba6b865a8cc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -50,7 +50,7 @@ use crate::execution::{ planner::operator_registry::OperatorRegistry, serde::{to_arrow_datatype, to_arrow_field}, }; -use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher}; +use crate::jvm_bridge::{jni_call, JVMClasses, OutputStreamSink, ShufflePartitionPusher}; use arrow::compute::CastOptions; use arrow::datatypes::{ DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit, DECIMAL128_MAX_PRECISION, @@ -352,6 +352,10 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + /// The output stream of the file the plan's Parquet writer writes, when the JVM created that + /// file, registered on the driving Spark task thread before native planning. Only a writer + /// planned to write through the JVM may use it. + output_stream: Option>, /// How comparisons treat floating-point operands. `Raw` only while planning a scan's data /// filters; see [`Self::create_data_filter`]. float_operands: FloatOperands, @@ -377,6 +381,7 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + output_stream: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, } @@ -500,6 +505,20 @@ impl PhysicalPlanner { self } + /// Attach the output stream of the file the plan's Parquet writer writes, for a file the + /// Spark task that created this plan created through the JVM. + /// + /// Only a writer planned to write through the JVM takes it, and such a writer requires it. + pub fn with_output_stream(mut self, output_stream: Option>) -> Self { + self.output_stream = output_stream; + self + } + + /// The output stream attached with [`Self::with_output_stream`], if any. + pub(crate) fn output_stream(&self) -> Option<&Arc> { + self.output_stream.as_ref() + } + /// How comparisons treat floating-point operands. pub fn float_operands(&self) -> FloatOperands { self.float_operands diff --git a/native/core/src/execution/planner/write.rs b/native/core/src/execution/planner/write.rs index 5665939afb2..29e28b8092e 100644 --- a/native/core/src/execution/planner/write.rs +++ b/native/core/src/execution/planner/write.rs @@ -32,6 +32,7 @@ use crate::{ spark_plan::SparkPlan, }, extract_op, + jvm_bridge::OutputStreamSink, }; /// Builder for native Iceberg writes. @@ -104,7 +105,7 @@ impl OperatorBuilder for ParquetWriterBuilder { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let parquet_writer = Arc::new(ParquetWriterExec::try_new( + let parquet_writer = ParquetWriterExec::try_new( Arc::clone(&child.native_plan), writer.output_path.clone(), codec, @@ -113,7 +114,13 @@ impl OperatorBuilder for ParquetWriterBuilder { (!writer.output_schema.is_empty()) .then(|| convert_spark_types_to_arrow_schema(&writer.output_schema)), object_store_options, - )?); + )?; + let parquet_writer = Arc::new( + match jvm_output_stream(writer.jvm_output_stream, planner.output_stream())? { + Some(output_stream) => parquet_writer.with_output_stream(output_stream), + None => parquet_writer, + }, + ); Ok(( scans, @@ -126,3 +133,65 @@ impl OperatorBuilder for ParquetWriterBuilder { )) } } + +/// The output stream a Parquet writer writes to, when the JVM created its file. +/// +/// A writer planned to write through the JVM requires the stream its own Spark task registered, +/// since the file it would otherwise create itself is not the one the JVM opened. A writer +/// planned to create its file must not be given one, which would mean the JVM opened a file that +/// nothing writes. +fn jvm_output_stream( + planned: bool, + registered: Option<&Arc>, +) -> Result>, ExecutionError> { + match (planned, registered) { + (true, Some(output_stream)) => Ok(Some(Arc::clone(output_stream))), + (true, None) => Err(ExecutionError::GeneralError( + "A Parquet writer planned to write through the JVM requires the output stream of \ + the file the JVM created" + .to_string(), + )), + (false, Some(_)) => Err(ExecutionError::GeneralError( + "A Parquet writer that creates its own file cannot use an output stream the JVM \ + created" + .to_string(), + )), + (false, None) => Ok(None), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct DiscardingSink; + + impl OutputStreamSink for DiscardingSink { + fn max_write_size(&self) -> usize { + 1 + } + + fn write(&self, _data: &[u8]) -> datafusion::common::Result<()> { + Ok(()) + } + } + + #[test] + fn a_writer_uses_the_output_stream_only_when_planned_to() { + let sink: Arc = Arc::new(DiscardingSink); + + let chosen = jvm_output_stream(true, Some(&sink)).unwrap(); + assert!(Arc::ptr_eq(&chosen.unwrap(), &sink)); + assert!(jvm_output_stream(false, None).unwrap().is_none()); + + let missing = jvm_output_stream(true, None).unwrap_err().to_string(); + assert!(missing.contains("requires the output stream"), "{missing}"); + let unexpected = jvm_output_stream(false, Some(&sink)) + .unwrap_err() + .to_string(); + assert!( + unexpected.contains("cannot use an output stream"), + "{unexpected}" + ); + } +} diff --git a/native/jni-bridge/src/lib.rs b/native/jni-bridge/src/lib.rs index 53228846591..9bf395353ac 100644 --- a/native/jni-bridge/src/lib.rs +++ b/native/jni-bridge/src/lib.rs @@ -196,6 +196,7 @@ mod comet_metric_node; mod comet_s3_credential_dispatcher; mod comet_task_memory_manager; mod comet_udf_bridge; +mod output_stream_sink; mod shuffle_block_iterator; mod shuffle_partition_pusher; @@ -204,6 +205,7 @@ pub use comet_metric_node::*; pub use comet_s3_credential_dispatcher::CometS3CredentialDispatcher; pub use comet_task_memory_manager::*; use comet_udf_bridge::CometUdfBridge; +pub use output_stream_sink::{JavaOutputStreamSink, OutputStreamSink}; use shuffle_block_iterator::CometShuffleBlockIterator; pub use shuffle_partition_pusher::{JavaShufflePartitionPusher, ShufflePartitionPusher}; diff --git a/native/jni-bridge/src/output_stream_sink.rs b/native/jni-bridge/src/output_stream_sink.rs new file mode 100644 index 00000000000..e7bbaff3b0b --- /dev/null +++ b/native/jni-bridge/src/output_stream_sink.rs @@ -0,0 +1,141 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Debug, Formatter}; + +use crate::{check_exception, errors::CometError, JVMClasses}; +use datafusion::common::{DataFusionError, Result}; +use jni::objects::{Global, JByteArray, JMethodID, JObject, JValue}; +use jni::signature::{Primitive, ReturnType}; +use jni::Env; + +/// Receives the bytes of a file that the JVM created, in order. +/// +/// The JVM creates the file, closes it once the native writer has written all of it, and aborts +/// it when the write fails, so an implementation only ever sees the file's bytes. One writer +/// writes one file, so calls never overlap. +pub trait OutputStreamSink: Send + Sync { + /// The most bytes one call to [`Self::write`] may carry. + fn max_write_size(&self) -> usize; + + /// Appends `data`, at most [`Self::max_write_size`] bytes, to the file. + fn write(&self, data: &[u8]) -> Result<()>; +} + +impl Debug for dyn OutputStreamSink { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("OutputStreamSink") + .field("max_write_size", &self.max_write_size()) + .finish() + } +} + +/// Writes to a JVM `org.apache.comet.parquet.HadoopFileOutput` from any native thread. +/// +/// Each write copies its bytes into a Java array this sink allocates once, then hands the array +/// to `write(byte[], int)`, so a write costs one copy and one upcall but no Java allocation. The +/// global references keep the Java object, its class and the array alive for as long as the +/// cached method ID is used. +pub struct JavaOutputStreamSink { + output: Global>, + buffer: Global>, + write_method: JMethodID, +} + +impl JavaOutputStreamSink { + /// Bytes per upcall, large enough that the upcall costs little next to the copy it carries. + pub const MAX_WRITE_SIZE: usize = 1024 * 1024; + + /// Captures `output` while running on an attached JVM thread. + pub fn try_new(env: &mut Env<'_>, output: &JObject<'_>) -> Result { + if output.is_null() { + return Err(DataFusionError::Execution( + "The output stream of a native write must not be null".to_string(), + )); + } + + let output_class = env.get_object_class(output).map_err(CometError::from)?; + let write_method = env + .get_method_id( + &output_class, + jni::jni_str!("write"), + jni::jni_sig!("([BI)V"), + ) + .map_err(CometError::from)?; + let buffer = env + .new_byte_array(Self::MAX_WRITE_SIZE) + .map_err(CometError::from)?; + + Ok(Self { + output: env.new_global_ref(output).map_err(CometError::from)?, + buffer: env.new_global_ref(&buffer).map_err(CometError::from)?, + write_method, + }) + } +} + +impl OutputStreamSink for JavaOutputStreamSink { + fn max_write_size(&self) -> usize { + Self::MAX_WRITE_SIZE + } + + fn write(&self, data: &[u8]) -> Result<()> { + if data.len() > Self::MAX_WRITE_SIZE { + return Err(DataFusionError::Internal(format!( + "A write of {} bytes exceeds the output stream's limit of {} bytes", + data.len(), + Self::MAX_WRITE_SIZE + ))); + } + // Lossless: the length is bounded by MAX_WRITE_SIZE. + let length = data.len() as i32; + + JVMClasses::with_env(|env| { + // SAFETY: `i8` and `u8` have the same size and alignment, and every bit pattern is + // valid for both. + let bytes = + unsafe { std::slice::from_raw_parts(data.as_ptr().cast::(), data.len()) }; + self.buffer + .set_region(env, 0, bytes) + .map_err(CometError::from)?; + + // SAFETY: `write_method` was resolved against the output object's class with this + // exact argument list and void return type. The global object reference keeps its + // defining class alive. + let result = unsafe { + env.call_method_unchecked( + self.output.as_obj(), + self.write_method, + ReturnType::Primitive(Primitive::Void), + &[ + JValue::Object(self.buffer.as_obj()).as_jni(), + JValue::Int(length).as_jni(), + ], + ) + }; + + // Inspect the pending exception before consuming the JNI result, so that its + // throwable survives the DataFusion error boundary. + if let Some(exception) = check_exception(env)? { + return Err(exception.into()); + } + + result.map_err(CometError::from)?; + Ok(()) + }) + } +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index f9292cb3f37..4f42e015397 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -995,7 +995,7 @@ message ParquetWriter { // 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. + // writer uses this path verbatim, and creates the file itself unless `jvm_output_stream` is set. string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; @@ -1013,6 +1013,12 @@ message ParquetWriter { // Catalyst's target schema, including top-level nullability and Parquet field IDs. // Nested collection field IDs are carried by the corresponding DataType messages. repeated SparkStructField output_schema = 9; + // Whether the JVM has already created `output_path` through Hadoop's FileSystem API and + // registered its output stream with the plan (Native.setOutputStream). The native writer then + // writes the file's bytes to that stream instead of opening the path itself, and the JVM closes + // the file afterwards. Set per task by CometWriteFilesExec for every file that is neither local + // nor on HDFS, such as the files S3A's magic committer needs S3A itself to create. + bool jvm_output_stream = 10; } enum AggregateMode { diff --git a/pom.xml b/pom.xml index 2a737c9c20c..29c11e4c0ad 100644 --- a/pom.xml +++ b/pom.xml @@ -542,6 +542,39 @@ under the License. + + + org.apache.spark + spark-hadoop-cloud_${scala.binary.version} + ${spark.version} + test + + + org.apache.hadoop.thirdparty + hadoop-shaded-guava + + + org.apache.hadoop + hadoop-annotations + + + javax.xml.bind + jaxb-api + + + commons-logging + commons-logging + + + org.apache.arrow + * + + + + org.testcontainers @@ -1136,6 +1169,11 @@ under the License. ${comet.abfs.parity.record} + + false diff --git a/spark/pom.xml b/spark/pom.xml index b6306449527..466ca29a403 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -111,6 +111,10 @@ under the License. spark-hadoop-cloud_${scala.binary.version} tests + + org.apache.spark + spark-hadoop-cloud_${scala.binary.version} +