andygrove commented on code in PR #5652:
URL: https://github.com/apache/datafusion-comet/pull/5652#discussion_r3970028747
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -74,8 +75,139 @@ use crate::cloud::s3::credential_bridge::AccessMode;
use crate::execution::operators::iceberg_common::load_file_io;
/// Builder chain instantiated once per task and handed to the partitioning
wrapper.
-type IcebergDataFileWriterBuilder =
- DataFileWriterBuilder<ParquetWriterBuilder, DefaultLocationGenerator,
DefaultFileNameGenerator>;
+type IcebergDataFileWriterBuilder = DataFileWriterBuilder<
+ ParquetWriterBuilder,
+ TrackingLocationGenerator,
+ DefaultFileNameGenerator,
+>;
+
+/// `DefaultLocationGenerator` that records every location it hands to a file
writer.
+///
+/// iceberg-rust's writers keep the `DataFile`s they have finalized private
until `close`, and
+/// have no abort hook, so when a task fails partway through there is no other
way to learn which
+/// files it created. The recorded locations let `delete_task_files` clean up
after a failure the
+/// way iceberg-java's `DataWriter.abort()` does.
+#[derive(Clone, Debug)]
+struct TrackingLocationGenerator {
+ inner: DefaultLocationGenerator,
+ locations: Arc<Mutex<Vec<String>>>,
+}
+
+impl TrackingLocationGenerator {
+ fn new(data_location: String) -> Self {
+ Self {
+ inner: DefaultLocationGenerator::with_data_location(data_location),
+ locations: Arc::new(Mutex::new(Vec::new())),
+ }
+ }
+
+ fn locations(&self) -> Vec<String> {
+ self.locations
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ .clone()
+ }
+}
+
+impl LocationGenerator for TrackingLocationGenerator {
+ fn generate_location(&self, partition_key: Option<&PartitionKey>,
file_name: &str) -> String {
+ let location = self.inner.generate_location(partition_key, file_name);
+ self.locations
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ .push(location.clone());
+ location
+ }
+}
+
+/// Deletes the tracked files if the write task is dropped before it finished.
+///
+/// A task can end without its future ever observing an error: when the
JVM-side input iterator
+/// throws, `executePlan` returns that error straight from the JNI batch pull
and the JVM then
+/// releases the plan, dropping this future mid-flight. The guard turns that
drop into the same
+/// cleanup the explicit error path performs. It stays armed until the task's
output batch has
+/// been handed to the JVM, which is the point where the JVM takes over
cleanup ownership.
+struct AbortOnDrop {
+ file_io: FileIO,
+ generator: TrackingLocationGenerator,
+ armed: bool,
+}
+
+impl AbortOnDrop {
+ /// Every location the task's writers were handed, in the order they were
generated.
+ fn locations(&self) -> Vec<String> {
+ self.generator.locations()
+ }
+
+ /// Give up ownership without deleting: the files are now someone else's
responsibility.
+ fn disarm(&mut self) {
+ self.armed = false;
+ }
+
+ /// Delete the tracked files, awaiting completion, and give up ownership.
Preferred over the
+ /// `Drop` path wherever the failure is observed inside the task's own
future, so the deletes
+ /// finish before the task reports its error rather than racing the
runtime's shutdown.
+ async fn abort(&mut self) {
+ self.armed = false;
+ delete_task_files(&self.file_io, self.generator.locations()).await;
Review Comment:
Fixed in 38976e3c8 — disarming now happens after the await, as you asked,
and as the earlier explicit error path did:
```rust
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
