fix: Delete data files left by a failed native Iceberg write task

[andygrove]

Which issue does this PR close?

Closes #5618.

Rationale for this change

When a native Iceberg write task fails partway through, the data files it had already finalized stay in the table's data location. They are invisible to readers and remove_orphan_files eventually reclaims them, but iceberg-java deletes them synchronously (DataWriter.abort() calls SparkCleanupUtil.deleteTaskFiles), so on spot-heavy or preemption-prone clusters the native path silently accumulates orphans that the JVM path does not. This is the first phase-2 item of the native Iceberg writes epic (#5649).

What changes are included in this PR?

The native writer records every location it hands to a file writer, and exactly one side owns deleting those files at any moment: the native writer until its output batch reaches the JVM, the JVM from then on. The handoff is explicit and does not depend on decoding the manifest.

Inside the native writer. iceberg-rust's writers keep the DataFiles they have finalized private until close and have no abort hook, so the task cannot ask a failed writer what it wrote. Instead, the RollingFileWriterBuilder is given a TrackingLocationGenerator, a wrapper around DefaultLocationGenerator that records every location it hands out. If the input stream, a write, or the final close fails, the task deletes every recorded location through its FileIO before propagating the original error.

That explicit path is not enough on its own, and the end-to-end test proved it: when the JVM-side input iterator throws (a UDF failure upstream of the write, the most common shape), executePlan returns the error straight from its JNI batch pull and the JVM releases the plan, so the write task's future is dropped without ever seeing an error. An AbortOnDrop guard covers that case. If it is dropped while still armed it deletes the tracked files, synchronously on a throwaway current-thread runtime when dropped from a plain JVM thread (releasePlan), or spawned onto the current runtime when dropped from inside one.

The guard is not disarmed when the writer closes: run_write_task returns it, still armed, alongside the DataFiles, and the outer task disarms it only once build_output_batch has produced the batch. encode_data_files_as_manifest and build_output_batch are both fallible and run after the files exist, so a failure there awaits the guard's abort (delete, then disarm) rather than leaving the files for nobody.

Across the JNI boundary. Once the batch reaches the JVM, the native side is out of the picture, but the locations must not be recoverable only by decoding the manifest — the manifest decode is itself a step that can fail, and does so under heap pressure (a valid 864 KB manifest from 4,096 files raised OutOfMemoryError inside Iceberg's Avro reader in a 128 MiB JVM). So the native operator emits the locations as a second Binary column next to the manifest, framed as a big-endian i32 count, then a big-endian i32 byte length and the UTF-8 bytes for each location. Explicit lengths rather than a separator, so a path is never re-interpreted whatever it contains.

CometIcebergWriteExec.doExecute registers a WrittenFileCleanup task failure listener before pulling the native payload, owning nothing at first, and drainNativePayload hands it the decoded locations before it copies the manifest bytes out of the off-heap batch. Everything after that point — the manifest Array[Byte] copy, the Avro decode, the metrics rebuild, TaskCommit construction, serialization — happens with the listener already owning the files. The decode is strict about consuming the whole column, so a framing divergence between the two sides fails loudly on every native write rather than silently handing cleanup a truncated list.

IcebergReflection.deleteFilesQuietly prefers SupportsBulkOperations.deleteFiles and falls back to FileIO.deleteFile(String) per path. Deletion is best-effort on both sides: failures are logged, never returned, so the task failure Spark reports is still the real one. The location that was open at the time of a failure is included; deleting a path that was never materialized is a no-op.

One window remains open and is worth stating plainly: if building the location list is itself what exhausts the heap, nothing deletes. That is a few thousand short strings against the Avro decoder's full DataFile objects plus metrics maps, so it is a much smaller target, but it is not zero.

The "Failure handling" section of iceberg-writes.md is updated to describe the ownership handoff and drop the reference to this issue.

How are these changes tested?

Rust:

• The tracking generator (layout unchanged, every location recorded).
• Cleanup against an in-memory FileIO, including a recorded location that was never written.
• The drop guard: a disarmed guard deletes nothing; an armed guard dropped outside a runtime deletes synchronously.
• run_write_task hands back an armed guard after a successful write whose locations match the files on disk, and aborting it removes them — the packaging-failure path.
• The locations framing round-trips through a decoder mirroring the JVM's, including a path with a newline and a non-ASCII path.
• The manifest round-trip test also asserts the locations column names exactly the files the manifest does.

JVM (CometIcebergWriteActionSuite):

• A failure-injection test writes a single-task source with a two-row Comet batch size into a table with a one-byte write.target-file-size-bytes, so the rolling writer finalizes a file per batch, and a UDF throws on the seventh row. A control run with the same source and settings first proves the writer rolls into several files; the failing run then asserts the write planned natively (CometIcebergWriteExec in the failed plan), no snapshot was created, the pre-existing data file is untouched, and no other parquet file remains under the table's data location.
• WrittenFileCleanup leaves the table alone before it has been handed any locations, and deletes them all once it has, through a real table's HadoopFileIO.
• The locations framing round-trips, an empty column decodes to nothing, and a trailing byte fails loudly.
• deleteFilesQuietly through a real table's HadoopFileIO: written files are removed, a nonexistent path is tolerated, and a second call is a no-op.

The full Iceberg suite set was run locally on the default Spark profile (253 tests, all passing), and the tree cross-compiles against spark-3.5/scala-2.12 and spark-4.0.

Every native write in the suite now decodes the locations column with the strict framing check, so the two sides' framing is exercised on each of them.

[unikdahal]

We disarm the cleanup guard here as soon as run_write_task completes, but the outer IcebergWriteExec::execute still performs fallible encode_data_files_as_manifest(...).await? and build_output_batch(...)? operations afterward. If either fails, the data files have already been written, this guard is disarmed, and the JVM never receives a manifest from which it could recover their locations. That leaves the task's files orphaned.

Could we keep cleanup ownership alive through manifest encoding/output-batch construction and disarm only once the native result is successfully materialized? One option would be for run_write_task to return the cleanup token alongside the DataFiles and let the outer task disarm it after packaging succeeds.

[andygrove]

Good catch, you're right — that window was real. run_write_task now returns the still-armed guard alongside the DataFiles, exactly as you suggested, and the outer task disarms it only after build_output_batch has produced the batch. If encode_data_files_as_manifest or build_output_batch fails in between, the task awaits abort_guard.abort() (delete-then-disarm) before propagating the error, so the deletes finish rather than racing the runtime teardown on the Drop path.

a_successful_write_returns_an_armed_guard_that_can_still_delete_its_files pins the new contract: a successful write hands back a guard that is still armed and whose recorded locations match the files on disk, and aborting it removes them.

[unikdahal]

By the time this listener is registered, both drainAvroPayload(batches) and decodeManifestToDataFiles(manifestBytes, specId) have already run. A failure in either operation happens after the native writer has successfully produced its data files, but before this listener owns cleanup. Since the native guard has also already been disarmed on successful writer close, neither side can remove those files.

This is especially important for decodeManifestToDataFiles: the cleanup paths are currently recoverable only by successfully decoding the same manifest whose decode may fail. I think cleanup ownership/locations need to cross this boundary independently of successful manifest decoding (or native cleanup needs to remain armed until the JVM acknowledges successful decode).

[andygrove]

Agreed, and the "recoverable only by successfully decoding the same manifest whose decode may fail" framing is the crux — that's the part I'd got wrong. Cleanup ownership no longer depends on the decode at all.

The native operator now emits the locations it wrote as a second Binary column next to the manifest, in a trivial framing (a big-endian count, then a length plus UTF-8 bytes per location) that the JVM walks with a ByteBuffer. doExecute registers the failure listener before pulling the native payload, owning nothing at first, and drainNativePayload hands it the locations off that column before it copies the manifest bytes out of the off-heap batch — so the decode, the metrics rebuild, TaskCommit, and serialization are all covered, and so is the manifest Array[Byte] copy itself, which for a large manifest is where the pressure actually shows up.

The residual window is now just building the location list, which is a few thousand short strings against the Avro decoder's full DataFile objects plus metrics maps. The decode is strict about consuming the whole column, so a framing divergence between the two sides fails loudly on every native write rather than silently handing cleanup a truncated list.

[sunchao]

[P2] Adding evidence to the existing JVM handoff thread: a valid 864,547-byte manifest from 4,096 real Parquet files decoded successfully with more heap available, but the same Iceberg 1.11 decoder failed under controlled heap pressure. In a 128 MiB standalone JVM, the low-headroom run had about 6.7 MiB free before decoding and raised OutOfMemoryError inside the actual Avro reader. All task files remained.

At this boundary, drainAvroPayload has already released the native plan, while the new failure listener is not yet registered. Could cleanup ownership remain armed until decoding and listener registration succeed? A component ownership model retaining the exact native guard across that decoder failure subsequently deleted all 4,096 files. This was isolated component validation, not a full Spark/JNI run. The decoder behavior predates this PR. The concern is the uncovered handoff in this cleanup fix, not a newly introduced decoder regression.

Current CI has 53 successful, 12 running and 7 skipped checks. I did not run the full Comet or Spark suites.

[andygrove]

Thanks for the controlled-heap repro — that's a much sharper statement of the problem than the thread it's attached to, and it pushed me off the design I had.

Rather than keep the native guard armed across the JNI handoff, I made the JVM's knowledge of the written files independent of the decoder. The native operator emits the locations as a second Binary column beside the manifest, framed as a big-endian count then a length plus UTF-8 bytes per location, which the JVM walks with a ByteBuffer. doExecute registers the failure listener before pulling the payload, owning nothing at first; drainNativePayload hands it the locations off that column before it copies the manifest bytes out of the off-heap batch. So your boundary is covered on both sides of it: the manifest Array[Byte] copy and the Avro decode both happen with the listener already owning the files.

I went this way rather than holding the native guard until the JVM acknowledges because the ack has to cross JNI, and the native guard's lifetime is tied to plan release — which happens at task end, after failure listeners run, so the two would double-delete and the disarm-on-success point would have to be a new JNI entry point. Reporting the locations makes the handoff a single volatile store with nothing fallible in between. Ownership is now explicit and non-overlapping: native until the output batch reaches the JVM (which now also covers manifest encoding and output-batch construction — see the sibling thread), JVM from then on.

The one thing I did not close is building the location list itself: if that allocation is what OOMs, nobody deletes. That's a few thousand short strings against the decoder's full DataFile objects plus metrics maps, so it is a much smaller target, but it isn't zero and I'd rather say so than claim otherwise.

Also rebased onto main to clear the conflict.

[sunchao]

Correctness

Re-reviewed 09b064e3 against 7f1e0018, including the changes since 9eddaafa. The previously demonstrated Avro-decoder handoff gap is addressed: the task failure listener is registered before consuming native output, and it receives the separately framed file locations before copying or decoding the manifest. The native packaging concern is also addressed in the control flow: run_write_task returns an armed guard, and packaging errors await deletion before returning the original error.

The maintained Spark 3.5 and 4.0 branches invoke failure listeners before completion listeners and preserve the original task exception. That matches the new ownership sequence through manifest decoding, metric reconstruction and commit-message construction. Cleanup uses individual locations generated for this task attempt; it does not scan or recursively delete a table directory. Partition and task-attempt IDs remain part of file names. One new [P2] remains in the failure-cleanup path: abort now disarms the guard before awaiting deletion. Dropping that pending cleanup leaves the remaining files without an owner. The previous code disarmed after deletion. The inline comment identifies the cancellation path and the small ordering correction.

Validation

Verified the actual CI checkout 614d9a5a, its exact base/head parents and its entire tree against this head. The native job passed 1,188 tests, with four skipped, including the armed-guard and location-framing cases. All four added JVM cleanup/framing tests passed on Linux Spark 3.5, Linux Spark 4.0 and macOS Spark 4.0. The failure-injection test requires native execution, checks that the original error survives and preserves a pre-existing committed file. Native artifact IDs and digests match the inspected producers and consumers.

The snapshot has 63 successful, nine skipped and one failed check. The Spark 3.4 expressions failure stopped at a Maven dependency download with HTTP 403, before those tests ran. A bounded local Rust component test compiled the exact guard and deletion helper with controlled FileIO/runtime test doubles: canceling between two deletions left one file; moving disarm after the await left zero. This establishes the guard cancellation behavior, not a complete Spark/JNI reproduction. I did not rerun full local suites or the historical heap-pressure experiment. The listener test exercises ownership directly; it does not inject an Avro failure through the complete Spark/JNI path. Allocation failure before the locations are acquired remains outside that demonstrated coverage, as acknowledged in the discussion. Maintained Spark 3.4/4.1 source branches were unavailable.

Performance

The new handoff adds a second binary payload and temporary copies proportional to the number and length of file paths, once per task. It avoids a second manifest decode or storage listing. Native cleanup deletes sequentially; JVM cleanup uses bulk deletion when available. I found no material new performance regression from the update; no throughput or cleanup-latency benchmark was run.

Design

Returning the armed guard makes responsibility for packaging failures explicit. The separate locations column lets the JVM acquire cleanup responsibility without depending on successful Avro decoding, while preserving the existing commit-message protocol. Cleanup remains best effort; the tests do not establish deletion under every cancellation, process-loss or resource-exhaustion scenario.

Abstraction & complexity

WrittenFileCleanup has a narrow task-local role, and the explicit length framing preserves paths containing separators, newlines and non-ASCII characters. These additions serve the ownership transfer directly. I found no new actionable abstraction or complexity concern.

[sunchao]

Correctness

[P2] Keep the abort guard armed until deletion finishes

Clearing armed before this await removes the cancellation fallback while deletion is still in progress. If the cleanup future is dropped after a delete yields, Drop now returns immediately and the remaining task files have no owner; no output batch has reached the JVM listener. This can be reached through executePlan: a pending stream poll still calls pull_input_batches, and a JVM input error exits JNI before releasePlan drops the stream. The previous explicit error path disarmed only after delete_task_files(...).await. Please move disarming after the await so teardown can still retry the tracked locations.

A bounded component test compiled the exact guard and delete helper with controlled FileIO/runtime test doubles. Canceling between two deletions left one file with this ordering; moving disarm after the await left zero. This verifies guard cancellation, not a complete Spark/JNI reproduction.

[andygrove]

Fixed in 38976e3 — disarming now happens after the await, as you asked, and as the earlier explicit error path did:

async fn abort(&mut self) {
    delete_task_files(&self.file_io, self.generator.locations()).await;
    self.armed = false;
}

Your reasoning about re-deletion holds too, which is what makes this safe rather than just safer: if the cancelled run already removed a file, the Drop retry logs a failed delete instead of failing, because delete_task_files is best effort by construction. So the only thing the old ordering bought was losing the fallback.

I put the reasoning on the method rather than in the commit, since the ordering looks arbitrary otherwise and is easy to "tidy" back.

On reproducing it in-tree. I tried to write this against the existing in-memory FileIO and could not, for a reason worth recording: the memory backend completes every delete inside a single poll. I measured it — one poll, zero pendings for a four-file abort — so there is no yield point at which to cancel, and any test written against it passes under both orderings.

cancelling_abort_keeps_the_guard_armed therefore uses the filesystem FileIO over a tempdir, where the deletes genuinely yield. It drives abort() with a no-op waker so the first yielding delete strands the future, drops it there, and then asserts:

1. guard.armed is still true;
2. at least one file survived — otherwise the test would pass vacuously, which was the trap in the memory-backed version;
3. dropping the still-armed guard removes the remainder.

Against the old ordering it fails on the first of those:

a cancelled abort must not have given up ownership of the remaining files

That matches your component test's result (one file left with the old ordering, zero with the new) from the other direction.

[sunchao]

Correctness

Re-reviewed 493ce7f98bd114da83998305c56bf2eff680b0a5 against 424c31aa79d13fddf743ffa29bae3c6f146e6c5e, including the increment from 09b064e3 and the overlapping rebase changes.

The remaining P2 finding is addressed. AbortOnDrop::abort now clears armed only after delete_task_files(...).await returns. Dropping that future while a deletion is pending therefore preserves the guard's cleanup ownership. Dropping the guard retries its tracked locations. The helper continues past individual deletion failures, including already-removed paths, so a partial first pass does not prevent cleanup of the remainder. The earlier packaging and JVM manifest-decode fixes remain intact.

The rebase wraps CometLocationGenerator and records its actual returned paths before file creation. Cleanup therefore follows the current partition-directory spelling without reconstructing paths, and retains the base's void-spec and clustered-input error handling. I rechecked the maintained Spark 3.5/4.0 task-listener sources: task-body failure invokes failure listeners before completion listeners and rethrows the original exception. The unchanged JVM listener is registered before native output is pulled and acquires locations before the manifest copy/decode. The already acknowledged allocation-failure limit before location adoption remains a best-effort limitation, not a newly demonstrated regression. No new or remaining P1/P2 issue was verified.

The completed Linux native CI job compiled merge 07af42f3, whose parents are this base/head and whose full tree equals the reviewed head: 1,264 tests passed, five skipped. This includes cancelling_abort_keeps_the_guard_armed, which cancels after a pending filesystem deletion, asserts the guard remains armed and files remain, then checks that dropping it removes the remainder. It also includes the existing armed-guard, manifest/framing and rebased writer tests. This is native test evidence, not a complete Spark/JNI cancellation run. At September 9, 17:46 UTC, CI has 43 successful checks, 7 skipped, 16 running and 1 queued, with no reported failure. Overall CI remains pending. I ran no local native/JVM test or benchmark and did not rerun the historical heap-pressure experiment. Maintained Spark 3.4/4.1/4.2 sources remain unavailable.

Performance

Moving the flag assignment adds no work to a successful write or an uninterrupted abort. Cancellation may retry locations already attempted, which is confined to failure cleanup and is necessary to retain ownership. Existing tracking and handoff costs still scale with file count and path bytes, not row count. This update makes no measured performance claim, and I found no additional performance issue.

Design

The fix restores the cancellation invariant inside abort: only completion of deletion disarms the guard, so dropping a suspended deletion leaves the guard responsible for cleanup. The successful-output handoff remains unchanged, with the previously discussed best-effort limits. Reusing the existing best-effort deletion helper avoids a second cleanup protocol. The filesystem regression test checks both retained ownership and eventual deletion. It does not claim cleanup survives process death, runtime shutdown or storage refusal.

Abstraction & complexity

The change adds no production abstraction or additional state. Keeping the ordering explanation beside abort makes the cancellation requirement visible where a future refactor could break it. The tracking wrapper still delegates path construction to the base implementation and only records the result, so it does not duplicate partition formatting logic.