unikdahal commented on code in PR #5652:
URL: https://github.com/apache/datafusion-comet/pull/5652#discussion_r3919637641
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala:
##########
@@ -183,6 +186,23 @@ case class CometIcebergWriteExec(
val decoded =
if (manifestBytes.isEmpty) new java.util.ArrayList[AnyRef]()
else IcebergReflection.decodeManifestToDataFiles(manifestBytes, specId)
+ // From here on the JVM knows which files iceberg-rust wrote. If this
task fails before its
+ // commit message reaches the committer (metrics rebuild, `TaskCommit`
construction,
+ // serialization), delete them the way iceberg-java's
`DataWriter.abort()` would; failures
+ // inside the native writer itself are cleaned up on the native side.
+ Option(TaskContext.get()).foreach { tc =>
+ tc.addTaskFailureListener(new TaskFailureListener {
Review Comment:
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).
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -357,19 +478,29 @@ async fn run_write_task(
// Build the field-id-decorated target schema once per task; every batch
is cast against it.
let target_schema =
Arc::new(iceberg::arrow::schema_to_arrow_schema(&iceberg_schema).map_err(iceberg_err)?);
- while let Some(batch) = input.try_next().await? {
- let decorated = decorate_batch_with_field_ids(batch, &target_schema)?;
+ let outcome = async move {
+ while let Some(batch) = input.try_next().await? {
+ let decorated = decorate_batch_with_field_ids(batch,
&target_schema)?;
+ let _timer = write_time.timer();
+ writer
+ .write(
+ decorated,
+ fanout_splitter.as_ref(),
+ clustered_splitter.as_ref(),
+ )
+ .await?;
+ }
let _timer = write_time.timer();
- writer
- .write(
- decorated,
- fanout_splitter.as_ref(),
- clustered_splitter.as_ref(),
- )
- .await?;
+ writer.close().await
}
- let _timer = write_time.timer();
- writer.close().await
+ .await;
+ // Whether the input stream, a write, or the close failed, every file this
attempt created is
+ // orphaned from here on: nothing will commit it, and the retry uses
attempt-unique names.
+ if outcome.is_err() {
+ delete_task_files(&file_io, location_generator.locations()).await;
+ }
+ abort_guard.disarm();
Review Comment:
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 `DataFile`s and let the outer task disarm it after
packaging succeeds.
--
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]
