From 5e042cf58b5a74724d35aed49386b9ed02eb5ecd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 27 Oct 2026 22:45:13 +0000 Subject: [PATCH] feat: support the S3A magic committer in native Parquet writes (#6940) * 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. * test: make the magic committer's failed-write check isolate the writer's abort S3A's job abort also aborts pending uploads under the destination, through core-default.xml's fs.s3a.committer.abort.pending.uploads=true, which hid a writer that never aborted its file. Turn that off in the test, use one task so job abort cannot race other tasks, and run each committer binding alone. Add a unit test that the output stream writer sends nothing after a failed write. * feat: support the S3A magic committer in Spark 3.4 and 3.5 native writes The Spark 3.x writer runs the session's commit protocol with the write's Hadoop configuration since #4746, so it reaches S3A's committers as Spark's writer does. Both writers now create task files through one helper, NativeWriteUtils.writeFile, which creates any file that is neither local nor on HDFS through Hadoop and closes or aborts it, and both serdes share one file system gate. The magic committer tests run on every supported Spark version; in-memory commit tracking is skipped on Hadoop 3.3.4, which lacks it. * refactor: simplify the native write path through the JVM The native writer now fails a task with the output stream's own error, which keeps the Java exception for the JNI boundary to rethrow as it was, so HadoopFileOutput no longer records failures and the JVM no longer swaps exceptions. Both writers write a task file through one helper, CometWriteFilesExec.writeFile, which sets the file and stream flag, runs the plan, and closes or aborts the file. Also: the pusher and output stream registrations share their rules, the writer takes an optional stream, local writes report their counted bytes, planning skips object store options and Hadoop configurations it does not need, the committer factory keys come from Hadoop's constants, and the tests share their setup and read S3A's success manifest through SuccessData. * address review: drop the unreachable bucket committer fallback and say once why S3A creates the file * address review: enrich write errors as Spark does, test the Hadoop route's failure paths Fail a native write's task with TASK_WRITE_FAILED as Spark's writer does, keeping fetch failures and fast-failing on an existing file. Admit subclasses of PathOutputCommitProtocol at planning. Build each task's plan once and create its iterator through CometExec, register the output stream through register_before_execution directly, and move the supported destinations to the Parquet page's Writing section. Test a file of several MiB, a failure partway through it, a failed close, a stream that cannot abort, a retried task and a path with special characters, and against MinIO, a failure after S3A uploaded a part and failed task and job commits. * test: compare what a failed job commit leaves pending with Spark's writer S3A's magic committer leaves the uploads of committed tasks pending when job commit fails, for Spark's own writer as well, so check that the native write leaves the same and nothing visible. Scope each pending upload check to its table, and abort what a test leaves behind. * address review: test the fast-fail and commit paths, nested and temporal columns, and the sink's guards Cover the branches of the Hadoop route that had no test: a task file that already exists, a file that exists at task commit (with and without fastFailOnFileFormatOutput), nested and temporal columns compared with Spark's writer through chunked writes and the abort path, an INSERT into a table of the V2 file source, and the output stream sink's oversize and null checks. * fix: handle the fallible execution context lookup in setOutputStream * test: write complex and temporal types through the Hadoop route * test: leave the non-ASCII character out of the output path when the JVM cannot encode it (cherry picked from commit f1121ca146856ce3443e958f21ecc4989c51b966) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + dev/ci/check-suites.py | 1 + .../contributor-guide/memory_management.md | 6 +- docs/source/user-guide/latest/operators.md | 4 +- docs/source/user-guide/latest/parquet.md | 8 + native/core/src/execution/jni_api.rs | 90 ++- .../src/execution/operators/parquet_writer.rs | 425 +++++++++++-- native/core/src/execution/planner.rs | 16 +- native/core/src/execution/planner/write.rs | 89 ++- native/jni-bridge/src/lib.rs | 2 + native/jni-bridge/src/output_stream_sink.rs | 199 ++++++ native/proto/src/proto/operator.proto | 9 +- pom.xml | 38 ++ spark/pom.xml | 4 + .../comet/parquet/HadoopFileOutput.java | 114 ++++ .../org/apache/comet/CometExecIterator.scala | 7 +- .../main/scala/org/apache/comet/Native.scala | 16 +- .../operator/CometDataWritingCommand.scala | 36 +- .../serde/operator/CometWriteFiles.scala | 15 +- .../serde/operator/NativeWriteUtils.scala | 130 ++++ .../sql/comet/CometNativeWriteExec.scala | 41 +- .../spark/sql/comet/CometWriteFilesExec.scala | 156 ++++- .../apache/spark/sql/comet/operators.scala | 11 +- .../comet/hadoop/fs/RecordingFileSystem.java | 211 +++++++ .../org/apache/comet/CometS3TestBase.scala | 2 +- .../CometParquetWriterRetrySuite.scala | 62 ++ .../parquet/CometParquetWriterSuite.scala | 572 +++++++++++++++++- .../parquet/CometParquetWriterTestBase.scala | 65 +- .../ParquetWriteWithMagicCommitterSuite.scala | 365 +++++++++++ .../CometExecIteratorLifecycleSuite.scala | 6 +- 31 files changed, 2538 insertions(+), 164 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/CometParquetWriterRetrySuite.scala create mode 100644 spark/src/test/scala/org/apache/comet/parquet/ParquetWriteWithMagicCommitterSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 19404eb85c3..b1300b2ced8 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.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 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index abd894ab0b1..317653f1c50 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.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 diff --git a/dev/ci/check-suites.py b/dev/ci/check-suites.py index 163f03d3741..96a49cf4900 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.csv.CsvReadFromS3Suite", # manual test suite "org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite", # manual test suite (loads libhdfs, see #5023) "org.apache.comet.IcebergReadFromS3Suite", # manual test suite diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 401387a5d1b..9bc57fb5f62 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -429,7 +429,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 diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 334d34673f2..93d39973b0e 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/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). | diff --git a/docs/source/user-guide/latest/parquet.md b/docs/source/user-guide/latest/parquet.md index 1c96006a145..a6f6af80e04 100644 --- a/docs/source/user-guide/latest/parquet.md +++ b/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..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 diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index c768d81f3f2..621de42c834 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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::{ @@ -39,7 +41,10 @@ use crate::{ serde::to_arrow_datatype, shuffle::spark_unsafe::row::{process_sorted_partitions, process_sorted_row_partition}, }, - jvm_bridge::{JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher}, + jvm_bridge::{ + JVMClasses, JavaOutputStreamSink, JavaShufflePartitionPusher, OutputStreamSink, + ShufflePartitionPusher, + }, }; use std::collections::HashSet; @@ -587,6 +592,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, } @@ -895,6 +904,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, })) } @@ -933,23 +943,72 @@ fn register_shuffle_partition_pusher( exec_context: &mut ExecutionContext, make_pusher: impl FnOnce() -> CometResult>, ) -> 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( + started: bool, + slot: &mut Option>, + what: &str, + make: impl FnOnce() -> CometResult>, +) -> 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 = + 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). @@ -1480,7 +1539,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(), diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 81fed5e5994..cb6f2964769 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,9 @@ enum ParquetWriter { Box, String, ), + /// Writer for a file the JVM created through Hadoop's FileSystem API (`HadoopFileOutput`), + /// which only supplies the file's bytes. + OutputStream(ArrowWriter), } impl ParquetWriter { @@ -154,13 +159,17 @@ impl ParquetWriter { Ok(()) } + ParquetWriter::OutputStream(writer) => writer + .write(batch) + .map_err(|e| writer.inner_mut().stream_failure_or(e)), } } /// 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 +177,21 @@ 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, which the local and + /// output stream writers count as they write. The HDFS writer uploads its file in pieces and + /// reports none. + async fn close(self) -> std::result::Result, parquet::errors::ParquetError> { match self { - ParquetWriter::LocalFile(writer) => { - writer.close()?; - Ok(()) + ParquetWriter::LocalFile(mut writer) => { + // `finish` rather than `close`, which would drop the count of bytes written. + writer.finish()?; + Ok(Some(writer.bytes_written() as u64)) } #[cfg(feature = "hdfs-opendal")] ParquetWriter::Remote( @@ -223,9 +238,114 @@ 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. + writer + .finish() + .map_err(|e| writer.inner_mut().stream_failure_or(e))?; + Ok(Some(writer.bytes_written() as u64)) + } + } + } +} + +/// The task's error for `e`, a Parquet write that failed. A failed write to the JVM output stream +/// comes back as the stream's own error, which keeps the Java exception it carries, so the JVM +/// rethrows that exception as it was. Any other failure is reported with `what` and its text. +fn write_error(e: parquet::errors::ParquetError, what: &str) -> DataFusionError { + match e { + parquet::errors::ParquetError::External(source) => { + match source.downcast::() { + Ok(stream_failure) => *stream_failure, + Err(source) => DataFusionError::Execution(format!( + "{what}: {}", + parquet::errors::ParquetError::External(source) + )), + } + } + e => DataFusionError::Execution(format!("{what}: {e}")), + } +} + +/// 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, + /// What the failed write returned, until [`Self::stream_failure_or`] reports it. Parquet only + /// sees the `io::Error` that `Write` allows, so this keeps the error itself: for the JNI sink + /// it holds the Java exception the stream threw, which the JVM then rethrows as it was. + failure: Option, +} + +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, + failure: None, + } + } + + /// The error to report for `e`, a write of this file that Parquet failed: the stream's own + /// error when a write to the stream caused it. + fn stream_failure_or( + &mut self, + e: parquet::errors::ParquetError, + ) -> parquet::errors::ParquetError { + match self.failure.take() { + Some(failure) => parquet::errors::ParquetError::External(Box::new(failure)), + None => e, + } + } + + /// 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; + let error = std::io::Error::other(e.to_string()); + self.failure = Some(e); + return Err(error); } + 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() } } @@ -247,6 +367,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 +406,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`, when there is one: the output stream of the file + /// the JVM created at `output_path`, which the writer then never creates itself. + pub fn with_output_stream(mut self, output_stream: Option>) -> Self { + self.output_stream = output_stream; + self + } + /// Create an Arrow writer based on the storage scheme /// /// # Arguments @@ -296,6 +427,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 +438,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 +607,18 @@ 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 => 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(), + )? + .with_output_stream(self.output_stream.clone()), + )), _ => Err(DataFusionError::Internal( "ParquetWriterExec requires exactly one child".to_string(), )), @@ -525,6 +675,7 @@ impl ExecutionPlan for ParquetWriterExec { props, runtime_env, &object_store_options, + self.output_stream.clone(), )?; // Clone schema for use in async closure @@ -560,25 +711,19 @@ impl ExecutionPlan for ParquetWriterExec { batch }; - writer.write(&renamed_batch).await.map_err(|e| { - DataFusionError::Execution(format!("Failed to write batch: {}", e)) - })?; + writer + .write(&renamed_batch) + .await + .map_err(|e| write_error(e, "Failed to write batch"))?; // The writer cannot spill, so a row group the pool will not grant fails the task. reservation.try_resize(writer.memory_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 = writer + .close() + .await + .map_err(|e| write_error(e, "Failed to close writer"))? + .unwrap_or(0) as i64; // Update metrics with write statistics files_written.add(1); @@ -854,20 +999,12 @@ mod tests { pool: &Arc, batches: Vec, ) -> 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 temp_dir = tempfile::tempdir()?; let file = temp_dir.path().join("part-00000.parquet"); - let writer = ParquetWriterExec::try_new( - input, + let writer = writer_of( + batches, format!("file://{}", file.display()), ParquetCompression::None, - 0, - column_names, - None, - HashMap::new(), )?; let pool: Arc = Arc::::clone(pool); let runtime = RuntimeEnvBuilder::new() @@ -917,6 +1054,211 @@ 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(()) + } + } + + /// A writer of `batches`, read from memory, to `output_path`. + fn writer_of( + batches: Vec, + output_path: String, + compression: ParquetCompression, + ) -> 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))); + ParquetWriterExec::try_new( + input, + output_path, + compression, + 0, + column_names, + None, + HashMap::new(), + ) + } + + /// 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 writer = Arc::new( + writer_of(batches, output_path, ParquetCompression::Snappy)? + .with_output_stream(output_stream), + ); + 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. + #[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 own error, unchanged, + /// which for the JNI sink carries the Java exception the JVM then rethrows as it was. + #[tokio::test] + async fn test_parquet_writer_fails_when_its_output_stream_fails() -> 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!( + matches!(&error, DataFusionError::Execution(message) + if message == "injected output stream failure"), + "{error:?}" + ); + assert_eq!( + sink.writes().len(), + 2, + "the writer kept writing to a stream that had failed" + ); + Ok(()) + } + + /// Once a write to the stream has failed, the writer sends nothing more: neither what a + /// caller writes next, nor what a flush, such as the one dropping the Parquet writer makes, + /// would send. A failed send also consumes none of the bytes it was handed, as `Write` + /// requires of an error. + #[test] + fn test_output_stream_writer_sends_nothing_after_a_failed_write() { + let sink = Arc::new(RecordingSink { + failing_write: Some(0), + ..RecordingSink::new(4) + }); + let mut writer = OutputStreamWriter::new(Arc::clone(&sink) as Arc); + + // A full buffer is only sent when the next bytes arrive. + writer.write_all(b"abcd").unwrap(); + assert!(sink.writes().is_empty()); + writer + .write(b"e") + .expect_err("the stream fails its first write"); + assert_eq!(sink.writes(), vec![b"abcd".to_vec()]); + assert_eq!(writer.buffer, b"abcd", "the failed write consumed bytes"); + + writer + .write(b"e") + .expect_err("the writer must not write to a failed stream"); + writer + .flush() + .expect_err("the writer must not flush to a failed stream"); + assert_eq!( + sink.writes().len(), + 1, + "the writer wrote to a failed stream" + ); + } + #[tokio::test] #[cfg(feature = "hdfs-opendal")] #[ignore = "This test requires a running HDFS cluster"] @@ -1068,6 +1410,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 08fc23d030c..8c240bc5dc2 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -54,7 +54,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, @@ -197,6 +197,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, @@ -222,6 +226,7 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + output_stream: None, float_operands: FloatOperands::Normalize, strict_timestamp_labels: false, } @@ -345,6 +350,15 @@ 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 + } + /// 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..8d1f03df9f4 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,16 +105,22 @@ impl OperatorBuilder for ParquetWriterBuilder { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let parquet_writer = Arc::new(ParquetWriterExec::try_new( - Arc::clone(&child.native_plan), - writer.output_path.clone(), - codec, - planner.partition(), - writer.column_names.clone(), - (!writer.output_schema.is_empty()) - .then(|| convert_spark_types_to_arrow_schema(&writer.output_schema)), - object_store_options, - )?); + let parquet_writer = Arc::new( + ParquetWriterExec::try_new( + Arc::clone(&child.native_plan), + writer.output_path.clone(), + codec, + planner.partition(), + writer.column_names.clone(), + (!writer.output_schema.is_empty()) + .then(|| convert_spark_types_to_arrow_schema(&writer.output_schema)), + object_store_options, + )? + .with_output_stream(jvm_output_stream( + writer.jvm_output_stream, + planner.output_stream.as_ref(), + )?), + ); 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..d88df799321 --- /dev/null +++ b/native/jni-bridge/src/output_stream_sink.rs @@ -0,0 +1,199 @@ +// 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; + + /// Fails a write that has no stream to go to. + fn check_not_null(output: &JObject<'_>) -> Result<()> { + if output.is_null() { + return Err(DataFusionError::Execution( + "The output stream of a native write must not be null".to_string(), + )); + } + Ok(()) + } + + /// Returns the length of a write of `length` bytes as the `int` the JVM takes, or fails a + /// write larger than [`Self::MAX_WRITE_SIZE`], which does not fit the reused Java array. + fn checked_write_length(length: usize) -> Result { + if length > Self::MAX_WRITE_SIZE { + return Err(DataFusionError::Internal(format!( + "A write of {} bytes exceeds the output stream's limit of {} bytes", + length, + Self::MAX_WRITE_SIZE + ))); + } + // Lossless: the length is bounded by MAX_WRITE_SIZE. + Ok(length as i32) + } + + /// Captures `output` while running on an attached JVM thread. + pub fn try_new(env: &mut Env<'_>, output: &JObject<'_>) -> Result { + Self::check_not_null(output)?; + + 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<()> { + let length = Self::checked_write_length(data.len())?; + + 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(()) + }) + } +} + +#[cfg(test)] +mod tests { + use super::JavaOutputStreamSink; + use datafusion::common::DataFusionError; + use jni::objects::JObject; + + #[test] + fn accepts_writes_up_to_the_limit() { + assert_eq!(JavaOutputStreamSink::checked_write_length(0).unwrap(), 0); + assert_eq!( + JavaOutputStreamSink::checked_write_length(JavaOutputStreamSink::MAX_WRITE_SIZE) + .unwrap(), + JavaOutputStreamSink::MAX_WRITE_SIZE as i32 + ); + } + + #[test] + fn rejects_writes_over_the_limit() { + for oversized_length in [ + JavaOutputStreamSink::MAX_WRITE_SIZE + 1, + i32::MAX as usize, + i32::MAX as usize + 1, + usize::MAX, + ] { + let error = JavaOutputStreamSink::checked_write_length(oversized_length).unwrap_err(); + assert!( + matches!(error, DataFusionError::Internal(_)), + "expected an Internal error, got {error:?}" + ); + let message = error.to_string(); + assert!(message.contains(&format!("A write of {oversized_length} bytes exceeds"))); + assert!(message.contains(&JavaOutputStreamSink::MAX_WRITE_SIZE.to_string())); + } + } + + #[test] + fn rejects_a_null_output_stream() { + let error = JavaOutputStreamSink::check_not_null(&JObject::null()).unwrap_err(); + assert!( + matches!(error, DataFusionError::Execution(_)), + "expected an Execution error, got {error:?}" + ); + assert!(error.to_string().contains("must not be null")); + } +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 8e720a8897c..6fee9b3bff3 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1027,7 +1027,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; @@ -1045,6 +1045,13 @@ 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.writeFile, which both write + // operators call, for every file that is neither local nor on HDFS. See HadoopFileOutput for + // why. + bool jvm_output_stream = 10; } enum AggregateMode { diff --git a/pom.xml b/pom.xml index 847b837ca87..477a5a4242b 100644 --- a/pom.xml +++ b/pom.xml @@ -545,6 +545,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 @@ -1139,6 +1172,11 @@ under the License. ${comet.abfs.parity.record} + + false diff --git a/spark/pom.xml b/spark/pom.xml index b5469ff40a5..baded236ffc 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} +