Skip to content

feat: support the S3A magic committer in native Parquet writes - #6940

Open
andygrove wants to merge 7 commits into
apache:mainfrom
andygrove:s3a-magic-committer-native-writes
Open

andygrove wants to merge 7 commits into
apache:mainfrom
andygrove:s3a-magic-committer-native-writes

Conversation

@andygrove

@andygrove andygrove commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #6929.

Rationale for this change

Native Parquet writes cannot run where S3A's magic committer commits the write, and Spark 4.1 makes it the default whenever spark-hadoop-cloud is on the classpath (SPARK-47618). The magic committer needs S3A itself to create each file. S3A turns a file created under __magic_job-<id>/ into a multipart upload to the file's final key, leaves the upload pending, and records it for task commit, either in a .pending file or, with fs.s3a.committer.magic.track.commits.in.memory.enabled, in memory. A writer that creates the object on its own bypasses all of that, and the job then succeeds without its data.

This takes the first approach from the issue: S3A creates the file, and the native writer only supplies its bytes. The magic committer's files never go through object_store, so this does not depend on #6355.

It covers Spark 3.4 and 3.5 as well. The issue limits this to Spark 4.0+ because the 3.x writer used to build its Hadoop contexts without the session's settings, but #4746 has since made it run the session's commit protocol with the write's Hadoop configuration, so it reaches S3A's committers the same way Spark's writer does.

What changes are included in this PR?

This changes both native writers: CometWriteFilesExec on Spark 4.0+ and CometNativeWriteExec on Spark 3.4 and 3.5.

Per task, in both writers (CometWriteFilesExec.writeFile):

  • The native writer still creates local and HDFS files itself. Any other file the commit protocol returns, such as the magic committer's __magic…/__base/part-… paths, is created by the JVM with the task's Hadoop configuration (HadoopFileOutput). It uses the same FileSystem.create call that parquet-mr's HadoopOutputFile.create makes for Spark's writer. The native writer streams the encoded file into that stream. The JVM then closes the stream, which is where S3A records the pending upload, before the write's statistics ask for the file's size.
  • When the write fails, the stream is aborted if it supports that (S3A's streams do), so a truncated file is never recorded as a pending upload. Other streams are closed, as Spark's writer closes its file on failure. The task fails with the exception the stream threw, not just its text: the native writer returns the stream's own error, and the JNI boundary rethrows the Java exception it carries as it was.
  • The route is chosen from the path the commit protocol actually returned, so a wrong guess at planning cannot misplace data: whichever committer runs, a file that is neither local nor on HDFS is created through Hadoop.

Native:

  • OutputStreamSink and JavaOutputStreamSink in the JNI bridge. The JVM registers the sink with the plan before execution through Native.setOutputStream, as it does for the remote shuffle pusher. Each upcall copies up to 1 MiB into a Java array the sink allocates once.
  • ParquetWriterExec gets an OutputStream writer: an ArrowWriter over a buffering Write adapter. It never opens output_path. After a failed write it sends nothing more, not even what dropping the writer flushes, and the task fails with that write's error unchanged. Its bytes_written metric reports what it wrote, as the local writer's now does too, rather than asking the file system.
  • ParquetWriter.jvm_output_stream in the proto. The planner rejects a writer planned to write to the stream when none is registered, and a registered stream for a writer that creates its own file.

Planning, in CometWriteFiles and CometDataWritingCommand (NativeWriteUtils.unsupportedFileSystem):

  • An S3 destination is admitted when the magic committer will commit it (NativeWriteUtils.usesS3AMagicCommitter). That means PathOutputCommitProtocol, or Spark's default protocol with BindingParquetOutputCommitter, and a committer factory, taken from mapreduce.outputcommitter.factory.class or else mapreduce.outputcommitter.factory.scheme.<scheme>, that is either MagicS3GuardCommitterFactory or S3ACommitterFactory with committer name magic. The name is resolved the way S3ACommitterFactory resolves it: the job's fs.s3a.committer.name first, then the bucket's.
  • Other S3 writes still fall back to Spark. Routing every S3A write through S3A would now be a one-line change to this gate, and would also cover the encryption, ACL, storage class and header settings feat: native Parquet writes to S3 #6355 falls back on. That is the open question in Support the S3A magic committer in native Parquet writes to S3 #6929, and it needs benchmarks against the object_store upload.

Tests and docs:

  • spark-hadoop-cloud's main jar is now in test scope, for PathOutputCommitProtocol and BindingParquetOutputCommitter. With it on the classpath, Spark 4.1 would enable the magic committer for every test session, so the scalatest system properties set spark.hadoop.fs.s3a.committer.magic.enabled=false, and the magic committer suite turns it back on.
  • operators.md documents the supported destinations.

Cost of the JVM route:

  • Encoding stays native, and only the encoded Parquet bytes cross into the JVM. Against the native writer writing a local file, each file costs two extra copies of the encoded bytes (into the writer's 1 MiB buffer, then into a Java array), one JNI upcall per MiB, 1 MiB of native memory counted against the task's memory pool, and one 1 MiB Java array allocated once per file.
  • Measured locally on 20M rows with a release build, through Hadoop's RawLocalFileSystem so that only the route differs, the route costs nothing beyond run-to-run noise: median 2172 ms against 2176 ms for snappy output (269 MiB), and 1926 ms against 1896 ms uncompressed (643 MiB, about 1% more CPU). GC time does not change. Spark's own writer takes about 5.8 s on the same data. The setup and full numbers are in Upload S3A magic committer files natively instead of through the JVM #6944.
  • The upload itself is S3A's and identical to Spark's writer. By default S3A stages each file on local disk in up to four 64 MiB blocks and uploads them with the AWS Java SDK. Upload S3A magic committer files natively instead of through the JVM #6944 tracks whether a native upload would beat that on real S3, and the options for doing it.

How are these changes tested?

  • Rust unit tests:
    • The stream writer's output is byte-identical to the local writer's, sent in writes no larger than the sink's limit. It never opens its path, here a magic committer path, and its bytes_written matches the file.
    • A failed write to the sink fails the task. Once a write has failed, the Write adapter sends nothing more, not even what a later flush would send, and a failed send consumes none of the bytes it was given.
    • The output stream can be registered only once, and only before execution.
    • The planner validates the stream against the proto flag.
  • CometParquetWriterSuite, which runs in CI:
    • A commit protocol hands out task files on a local-disk test FileSystem under its own scheme, which the native writer cannot open. The write runs natively, and every file is created through Hadoop and then closed.
    • A failing stream gets its file aborted, and the task fails with the stream's own IOException.
    • The planning check resolves the committer correctly across bindings, factory keys, committer names, and per-bucket settings.
  • ParquetWriteWithMagicCommitterSuite, new and manual like the other MinIO suites, listed in dev/ci/check-suites.py. It uses the magic committer as Spark itself configures it, by default on 4.1+ and through the per-bucket key on earlier versions:
    • The _SUCCESS manifest names the magic committer, no __magic directory or pending upload is left, and the write's reported byte count matches the objects. On 4.0+ that count comes from the marker S3A writes when it closes a magic file, so the check also covers the close ordering. On 3.x it comes from the native writer's own count.
    • The same with commits tracked in memory, where S3A has that (not on Hadoop 3.3.4, which Spark 3.4 and 3.5 ship), and through each binding alone.
    • A failed single-task write leaves no data file and no pending upload. The test turns off fs.s3a.committer.abort.pending.uploads, which core-default.xml turns on, because S3A's job abort would otherwise abort the upload itself and hide a writer that never aborts its file.
    • Writes under the file and directory committers stay on Spark.

Run locally on macOS (arm64):

  • cargo test -p datafusion-comet --lib, filtered to the writer and registration tests: 12 passed, and the 4 HDFS tests are ignored as before.
  • CometParquetWriterSuite plus ParquetWriteWithMagicCommitterSuite:
    • Spark 4.0, 4.1 and 4.2: 73 passed on each.
    • Spark 3.5: 69 passed and 4 canceled, which are three tests that need Spark 4.0+ and the in-memory tracking test.
    • Spark 3.4: 67 passed and 6 canceled, the same four plus two tests that are skipped on 3.4 already.
  • The magic committer suite ran against S3Proxy 2.9.0 in place of MinIO, since there is no Docker here, through a local-only CometS3TestBase patch that is not part of this PR.
  • Mutations, each caught by at least one of the tests above: not aborting the file on failure (on 3.5 and 4.1), not rethrowing the stream's exception, not closing the file after a successful write, planning ignoring the session's commit protocol, a flush that sends nothing, and a writer that forgets a failed write.
  • cargo clippy --all-targets --workspace -- -D warnings on stable and on nightly 1.99, cargo fmt, spotless, scalastyle, scalafix in check mode on 3.5, and the strict-warnings compile on 3.5 with the delta contrib.

Not run: the magic committer suite against real MinIO.

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.
@github-actions github-actions Bot added enhancement New feature or request area:writer Native Parquet writer area:scan Parquet scan / data reading area:ffi Arrow FFI / JNI boundary labels Oct 13, 2026
…r'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.
@andygrove

Copy link
Copy Markdown
Member Author

@parthchandra @comphead I'm still iterating om this, but let me know what you think of the direction

The Spark 3.x writer runs the session's commit protocol with the write's
Hadoop configuration since apache#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.
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.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the full four-commit, 24-file diff from 474b5ece841a7fc6dff37ae04bd6787e654d7be7 to 4756615e289db86dda118aec72d12b87e638a7ee. The PR is not a draft. Read the supplied discussions and linked issues. No introduced P1/P2 issues found within this review, and no substantiated existing P1/P2 blockers remain.

Routed skills: review-comet-pr, review-comet-expression-pr for serde/planning, review-comet-ffi-pr, review-comet-memory-pr, and review-comet-shuffle-pr for the shared callback-registration refactor.

Summary

  • Prior state and problem: Native Parquet writes admitted local and HDFS destinations but fell back for S3. Writing magic-committer paths directly would bypass S3A’s pending-upload tracking and could leave a successful job without committed data.
  • Design approach: Hadoop creates the task file and owns its upload lifecycle. Native code supplies encoded Parquet bytes through a registered JNI sink. Planning admits recognized magic-committer configurations, while execution selects the transport from the actual task filename.
  • Correctness: Checked Spark’s commit protocols, file-writer cleanup and statistics ordering, Hadoop’s committer selection and stream lifecycle, and parquet-rs’s footer flushing. The stream closes before statistics and task commit. Failed native writes abort the Hadoop stream, and stream exceptions retain their Java throwable through the native boundary. No introduced P1/P2 correctness issue was identified.
  • Compatibility analysis: Compared relevant sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0, plus Hadoop 3.3.4 and 3.4.1. The factory and committer-name precedence matches Hadoop. Both Spark writer entry points use the shared transport, without introducing Spark-version logic into Rust. Native writes remain experimental and disabled by default.
  • Key design decisions: S3A retains responsibility for magic paths, pending commits, encryption and upload settings. The per-task jvm_output_stream flag is checked against the registered sink, preventing a missing or unexpected stream from silently selecting another destination.
  • Implementation sketch: HadoopFileOutput handles create/write/close/abort. JavaOutputStreamSink reuses a Java byte array. OutputStreamWriter batches writes and stops sending after a sink failure. CometWriteFilesExec.writeFile shares execution and cleanup across Spark 3.x and 4.x. Local byte metrics now use the writer’s byte count.
  • Performance: The new route adds two copies of encoded bytes, approximately one JNI upcall per MiB, a 1 MiB native buffer included in writer accounting, and a reusable 1 MiB Java array. The author’s measurements in #6944 show little local routing overhead, but do not establish real-S3 upload performance. No evidence-backed regression was identified on the existing local/HDFS paths.
  • Design: Keeping S3A responsible for its own commit protocol is straightforward and avoids duplicating Hadoop-internal metadata formats. Choosing transport from the committer’s returned path also makes execution resilient to planning-time assumptions.
  • Abstraction & complexity: The sink interface provides a focused boundary between encoding and transport. Shared file execution removes duplication between Spark versions. The extracted callback-registration helper preserves remote shuffle’s existing once-only, before-execution checks, with its regression test passing at this head.
  • Behavioral changes worth calling out: Compared with branch-1.1 at 8ac6fccde53d96fec22184be287747eeae481cac, eligible magic-committer writes can now use native encoding when native writes are enabled. Other S3 configurations still fall back. Earlier Spark 3.x commit-lifecycle changes already present in the base were distinguished from this PR’s changes. No unintended release-relative regression was identified.
  • Suggested improvements: No additional implementation change met the P1/P2 reporting bar. Broader Spark SQL and object-store validation remain the limitations described below.

Exact-head CI: run 38475818148 has 26 successful and 18 skipped checks, with no failed checks. Logs confirm 2,429 Rust tests passed, including the new stream and planner tests. The Spark 4.1 scans job passed the new Hadoop-routing, abort/exception and committer-selection tests. Spark SQL jobs across supported versions, Iceberg integration jobs, macOS and benchmarks were skipped.

Local validation: compiled the exact HadoopFileOutput and RecordingFileSystem sources with JDK 17, cached Hadoop 3.4.2 and parquet-java 1.16.0. A disposable probe passed byte preservation, escaped/Unicode paths, explicit close, no-overwrite behavior, original IOException propagation and abort cleanup. Suite registration and git diff --check also passed. Project code remains unchanged.

Validation limits: local Rust tests stopped before compilation because the configured registry could not resolve locked DataFusion 55.2.0, both offline and online. Full local Comet integration suites were therefore not rerun. I did not independently run MinIO, real S3, or the multi-version magic-committer suite. The author’s reported S3Proxy runs across all five Spark profiles are supplementary evidence, not independently reproduced results.

@andygrove
andygrove requested a review from viirya October 17, 2026 14:08

@viirya viirya left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this. Having S3A create the file and letting the native writer supply only the bytes keeps the magic committer's protocol entirely inside Hadoop. Choosing the route per task from the committer's actual path is a nice property. I traced the upcalls and they run on the Spark task thread through block_on, since the writer's plan always has a JVM-fed scan, so S3A blocking on its upload queue never ties up a tokio worker. The close-before-closeFile ordering and the abort path also look right to me. I left a few inline comments on the S3 gate and the comments around it. Two more things don't anchor to a line in the diff.

Validation. The end-to-end evidence for this PR comes from ParquetWriteWithMagicCommitterSuite. As committed, that suite runs against MinIO through CometS3TestBase, but the runs so far used S3Proxy through a local patch, so the code that would merge hasn't run yet. Could you run it once as committed on 3.5 and 4.1? If Docker isn't available, making the backend selectable in CometS3TestBase, so the S3Proxy setup is reproducible from the PR itself, would work too. This PR also changes how every native write's task file is routed, so could you trigger the manual Spark SQL Writer Tests workflow for 3.5 and 4.1?

Memory docs. The ParquetWriterExec bullet under "Native writers" in memory_management.md should mention the 1 MiB output-stream buffer that the reservation now includes. It should also say that S3A's own upload buffers (fs.s3a.fast.upload.buffer) sit outside Comet's pool, as they do for Spark's writer.

* `None` when it can. Local and HDFS files are created natively. A write that S3A's magic
* committer commits is supported too, since each of its files is created through S3A (see
* `CometWriteFilesExec.writeFile`), which is the only way to give the committer the pending
* upload it completes. Other S3 writes stay on Spark until the native writer reproduces what

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The per-task routing makes every S3A-served write correct. Staging committers hand out local file: paths, which the native writer creates itself, and the file committer's _temporary paths are created through Hadoop. So this gate now only decides scope, and "until the native writer reproduces what S3A does at create time" no longer describes why other S3 writes fall back. The same reason appears in the last test of ParquetWriteWithMagicCommitterSuite.

Since #6355 hasn't landed, the alternative for a non-magic S3A write today is Spark's writer, not object_store. Would admitting every S3A destination be the simpler choice? It would drop usesS3AMagicCommitter and its test, which mirror Hadoop's committer resolution and could drift as Hadoop changes. If you'd rather keep the scope, could the comments here and in the suite say the real reason (benchmarks against #6355)?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am keeping the scope. Non-magic S3A writes stay on Spark until streaming a file through S3A has been benchmarked against uploading it natively (#6944), so the reason is scope, not correctness, and your point that the old wording no longer describes it is right. The doc comment on unsupportedFileSystem and the last test of ParquetWriteWithMagicCommitterSuite now say that. (65e2bfd)

factory match {
case Some(MagicCommitterFactoryClass) => true
case Some(S3ACommitterFactoryClass) =>
// S3ACommitterFactory takes the name from the job's configuration when it has one, and

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

core-default.xml sets fs.s3a.committer.name=file in both 3.3.4 and 3.4.1, so with a defaults-loaded configuration this lookup always finds a job-level name, and the bucket fallback never applies. That matches S3ACommitterFactory, which also reads the task conf first, so it isn't wrong. But Spark never reaches this branch. Could it be dropped?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropped. Spark's job configuration always has a name from core-default.xml, and S3ACommitterFactory reads the job's first, so a bucket's own name never decides it. The doc comment of usesS3AMagicCommitter says so. (65e2bfd)

!usesMagic(pathOutput, s3aSchemeFactory, "fs.s3a.committer.name" -> name)(),
s"committer name '$name'")
}
// A bucket's committer name applies when the job sets none of its own, and a global name in

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These cases build new Configuration(false), so "a bucket's committer name applies when the job sets none" describes a setup Spark never has. Spark's Hadoop conf carries fs.s3a.committer.name=file from core-default.xml and, once JobConf has loaded, the s3a scheme factory from mapred-default.xml. Could you add a case with new Configuration() that pins the behavior Spark actually sees, or drop these assertions together with the fallback?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Replaced. The Configuration(false) bucket assertions are gone with the fallback. A new Configuration() case now asserts that the default name is file, that a bucket's magic does not win over it, and that a global magic name does, including against a bucket's file. (65e2bfd)

Comment thread native/proto/src/proto/operator.proto Outdated
// 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 and CometNativeWriteExec for every

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only CometWriteFilesExec.writeFile sets this flag. CometNativeWriteExec reaches it through that helper. More generally, the explanation of why the magic committer needs S3A to create the file now appears in about eight places: here, HadoopFileOutput, both exec scaladocs, writeFile, NativeWriteUtils, parquet_writer.rs and operators.md. Could it live once, in HadoopFileOutput or writeFile, with the other places referring to it? This sentence shows the copies are already drifting.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. The full explanation is now in the HadoopFileOutput class doc, and the other places refer to it: both exec scaladocs, writeFile, the NativeWriteUtils docs, parquet_writer.rs, and the test comments. operators.md keeps only the user-facing sentence. The proto comment now says that CometWriteFilesExec.writeFile sets the flag, and that both write operators call it. (65e2bfd)

@andygrove

Copy link
Copy Markdown
Member Author

Thanks @viirya. The memory_management.md bullet now mentions the 1 MiB output-stream buffer in the reservation, and says S3A's own upload buffers (fs.s3a.fast.upload.buffer) sit outside Comet's pool. (65e2bfd)

On validation: I have no Docker on my machine, so I cannot run ParquetWriteWithMagicCommitterSuite against MinIO here. I am running it as committed on Spark 3.5 and 4.1 on a Docker host, and dispatching the manual Spark SQL Writer Tests workflow for 3.5 and 4.1, and I will post the results here.

@viirya viirya left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the quick turnaround. The scope wording, the dropped bucket fallback, the new Configuration() cases, the consolidated docs and the memory note all look good to me. I'll wait for the MinIO and Spark SQL Writer Tests results.

One thing I meant to raise last time and dropped: nativeWriterCreatesFile returns false for every s3a: path. So if #6355 lands after this, its S3 writes would be routed through HadoopFileOutput instead of object_store, and its tests would still pass because the data gets written either way. cc @comphead: could #6355's tests assert which route a write took, or could that PR note that nativeWriterCreatesFile needs to change with it? The jvm_output_stream check covers data loss, but not this.


/**
* A fallback reason when the native writer cannot write to `outputPath`'s file system, or
* `None` when it can. Local and HDFS files are created natively, and so are the files of a

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small wording slip from the reshuffle: this now says the magic committer's files are "created natively ... through S3A". The point is the opposite, they're created by S3A and only written natively. Maybe "Local and HDFS files are created natively. The files of a write that S3A's magic committer commits are created through S3A and written natively (see HadoopFileOutput)."

@comphead

Copy link
Copy Markdown
Contributor

I checked this on the cluster and it seems to work with magic and regular committer on Spark 4

@comphead comphead left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having S3A create the file while the encoding stays native looks right to me, and viirya has already covered the scope, the committer resolution and the doc consolidation, so I have not repeated those. Three comments on code that the refactor left behind and on where the user docs live now. The PR description also still lists a Rust test showing that the output stream "can be registered only once, and only before execution", but 4756615e28 removed register_output_stream_once_before_execution. The description becomes the squash commit message, so it needs an update.

Comment on lines +1000 to +1013
/// 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<Arc<dyn OutputStreamSink>>,
) -> CometResult<()> {
register_before_execution(
exec_context.root_op.is_some(),
&mut exec_context.output_stream,
"A native write's output stream",
make_sink,
)
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

register_output_stream now has one caller, setOutputStream, and no test of its own since 4756615e28 replaced register_output_stream_once_before_execution with the shared register_before_execution. register_shuffle_partition_pusher is still worth keeping as a function because its test calls it. Could setOutputStream call register_before_execution directly, and could the module doc at the top drop register_output_stream from its list of core functions? The once-and-before-execution rules stay covered by the pusher test, which goes through the same helper.

Comment on lines +352 to +369
private[comet] def writeFile(
filePath: String,
hadoopConf: Configuration,
nativeOp: Operator,
inputs: Array[Object],
numOutputCols: Int,
nativeMetrics: CometMetricNode,
numPartitions: Int,
partitionId: Int): Unit = {
def run(output: Option[HadoopFileOutput]): Unit = {
val writer = nativeOp.getParquetWriter.toBuilder
.setOutputPath(filePath)
.setJvmOutputStream(output.isDefined)
.build()
val execIterator = new CometExecIterator(
CometExec.newIterId,
inputs,
numOutputCols,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both callers build a complete Operator only for writeFile to take it apart again. writeNatively and CometNativeWriteExec each run getParquetWriter.toBuilder ... build() and then toBuilder.setParquetWriter(...).build(), and writeFile repeats both steps to set the path and the flag. Could writeFile take the template Operator together with the caller's ParquetWriter.Builder, and set output_path and jvm_output_stream on that builder? Each task would then build its plan once, and "set up for the task except for the file" would show in the parameter types rather than in the doc comment.

This helper also constructs CometExecIterator by hand, where both callers used CometExec.getCometIterator before. Adding outputStream: Option[HadoopFileOutput] = None to the existing overload would keep iterator construction in CometExec.

Comment on lines +143 to +149
Native Parquet writes support 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since #7090, parquet.md has a "Writing" section that holds the native writer's scope (non-partitioned and non-bucketed writes only, and not with spark.sql.files.maxRecordsPerFile), and the scan rows in this table link to parquet.md#reading. Could this paragraph move into that section, next to the scope sentence, with the WriteFilesExec and DataWritingCommandExec rows linking to parquet.md#writing? Someone checking which destinations the writer supports would then find them with the rest of its limits.

@parthchandra parthchandra left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This approach is excellent. The code looks good - I looked at the failure paths: task failures, retries, fetch failures, write and close failures, and task and job commit failures. They do seem toe be covered.
The only comments I have are about the tests -

  • Can we enable run-all-spark-profiles so all profiles run thru ci?
  • dev/ci/check-suites.py is the nly suite that runs s3a magic committer code and it is excluded from the ci. .github/workflows/delta_contrib_test.yml has a job (contrib-delta-s3) that runs a MinIO suite on a plain runner. Could we do the same for this suite? This can be a follow-up

// Hadoop's FileSystem API (see HadoopFileOutput) and the native writer only supplies the
// bytes. RecordingFileSystemCommitProtocol puts task files on a scheme the native writer
// cannot open at all, so the write can only succeed that way.
withRecordingFileSystem(rows = 1000, partitions = 3) { (df, outputPath) =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This writes rather small files. The 1MB chunking, the reused java array only get exercised when the file is bigger tahn 1 MB.
Could we add a test with a file of a few MiB? For example, use many rows of random non-ASCII strings with compression=none, and check that the data reads back. It would also help to let RecordingFileSystem fail the Nth write.

}

@Override
public void close() throws IOException {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add a failsClose switch and a test? (With S3A, the last part upload and the .pending. files are written in close() so this might be where we will likely get most failures)
We would expect the task to fail, abortTask to run and no part files to be left behind.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, hasCapability always has the stream as abortable. Can we also have a switch to turn that off?

}
}

test("a failed write through Hadoop aborts the file and fails with what its stream threw") {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we add a test for a task that is retried on this route? For example, the first attempt fails after its file is created and the second attempt succeeds. We would expect one aborted file, one closed file, each row once in the output, and nothing left over.

}
}

test("a failed native write leaves no pending upload behind") {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test fails in the child projection on the first batch, so the file is created but no bytes reach the stream. It covers an upload that was started and never fed. Could you add cases that show rollback with real S3A?

First, a source big enough to upload at least one part before the failure. For example, set fs.s3a.multipart.size to 5M and write more than 5 MiB. Then the abort has to cancel a part that was already uploaded.

Second, a commit protocol based on PathOutputCommitProtocol that throws in commitTask before calling super. The file is already closed in that case, so only S3A's abortTask can clean up.

Third, the same for commitJob, which on 3.x is Comet's own abortJob call.

After each case I would expect no visible data, no __magic directory and no pending upload.

private def withRecordingFileSystem(rows: Int, partitions: Int)(
write: (DataFrame, String) => Unit): Unit =
withTempPath { dir =>
val outputPath = new File(dir, "output.parquet").getAbsolutePath

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could the output directory in withRecordingFileSystem have a name with a space, %, #, ? and a non-ASCII character? The old object_store route had to decline paths like these because the JVM and native sides spelled them differently. This route builds a Hadoop Path from the string the committer returned, so I expect it to work. A test would show it and would stop a later change to HadoopFileOutput.create from breaking it.

if (NativeWriteUtils.nativeWriterCreatesFile(filePath)) {
run(None)
} else {
val file = HadoopFileOutput.create(filePath, hadoopConf)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Spark's writer wraps a failed task write in TASK_WRITE_FAILED ("Task failed while writing rows to "). It lets FetchFailedException through. In the default setup it turns FileAlreadyExistsException into TaskOutputFileAlreadyExistException, which tells Spark not to retry the task. See FileFormatDataWriter.enrichWriteError, and FileFormatWriter.executeTask on 3.x.

executeTask here rethrows the raw error. This route makes the gap easier to hit. HadoopFileOutput.create uses overwrite = false, so it can now throw FileAlreadyExistsException. Spark would stop the job right away, and Comet will retry the task up to spark.task.maxFailures times.

The test comment at CometParquetWriterSuite.scala:2057 says the failure surfaces as it would in Spark's writer. That is not true for the error class. Is this intended? If it is out of scope here, could you open an issue and link it from executeTask? Otherwise a small wrapper with the same three cases would match Spark.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ffi Arrow FFI / JNI boundary area:scan Parquet scan / data reading area:writer Native Parquet writer enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support the S3A magic committer in native Parquet writes to S3

5 participants