sunchao commented on code in PR #5652:
URL: https://github.com/apache/datafusion-comet/pull/5652#discussion_r3954175319
##########
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:
### 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.
--
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]
