This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 8320ae481b fix: Delete completed tasks' data files when an Iceberg
write job fails (#5663)
8320ae481b is described below
commit 8320ae481b4eb916beebc9ed611473abb79c3e71
Author: Andy Grove <[email protected]>
AuthorDate: Thu Sep 10 21:54:42 2026 -0600
fix: Delete completed tasks' data files when an Iceberg write job fails
(#5663)
* fix: delete data files left by a failed native Iceberg write task
iceberg-java's writer abort deletes the files a failed task attempt
wrote; the native path left them for remove_orphan_files. Close the gap
in both places a task can fail.
Inside the native writer, a TrackingLocationGenerator records every
location handed to a file writer, since iceberg-rust's writers keep
finalized files private until close and have no abort hook. The task
deletes the recorded locations when a write fails, and an AbortOnDrop
guard does the same when the task future is dropped without ever seeing
an error, which is what happens when the JVM input iterator throws:
executePlan returns that error from its JNI batch pull and the JVM
releases the plan.
After the native writer has returned, CometIcebergWriteExec registers a
task failure listener that deletes the decoded manifest's files through
the table FileIO, via a new best-effort IcebergReflection helper.
Both deletions log failures rather than raising them, so the original
task failure is the one Spark reports.
Closes #5618
* fix: keep the Iceberg abort guard armed until deletion finishes
* fix: delete completed tasks' data files when an Iceberg write job fails
IcebergCommitExec collected task commit messages with executeCollect,
which only returns once every task has succeeded, so a job failure left
the committer with no messages and the data files of the tasks that had
completed stayed in the table's data location.
Collect each task's message as it finishes through a runJob result
handler, as Spark's own V2 write does, and abort with the completed
messages. Iceberg's SparkWrite.abort only deletes files after a
cleanable commit failure and skips cleanup before any commit was
attempted, so the committer then deletes the completed tasks' data files
itself through the table FileIO; nothing can reference them at that
point. The failing task's own files are handled by task-level cleanup.
Closes #5277
---------
Co-authored-by: test <[email protected]>
---
docs/source/user-guide/latest/iceberg-writes.md | 33 +-
.../core/src/execution/operators/iceberg_write.rs | 604 +++++++++++++++++++--
.../apache/comet/iceberg/IcebergReflection.scala | 60 ++
.../spark/sql/comet/CometIcebergWriteExec.scala | 105 +++-
.../apache/spark/sql/comet/IcebergCommitExec.scala | 74 ++-
.../org/apache/comet/CometIcebergTestBase.scala | 26 +
.../comet/CometIcebergWriteActionSuite.scala | 298 +++++++++-
7 files changed, 1109 insertions(+), 91 deletions(-)
diff --git a/docs/source/user-guide/latest/iceberg-writes.md
b/docs/source/user-guide/latest/iceberg-writes.md
index 20b9635f59..e179606a9b 100644
--- a/docs/source/user-guide/latest/iceberg-writes.md
+++ b/docs/source/user-guide/latest/iceberg-writes.md
@@ -209,17 +209,28 @@ attempt id is embedded in its data file names.
Partial results are never committed. The commit set is exactly the commit
messages returned by
successful tasks — a failed task contributes none — and if the job fails, the
driver-side
-commit operator aborts without committing anything. Data files already
finalized by a failed
-task attempt are not deleted by that task (iceberg-java's writer abort deletes
them; the
-native path has no abort hook yet — tracked in
-[#5618](https://github.com/apache/datafusion-comet/issues/5618)): they are
invisible to every
-reader, since readers resolve files through committed manifests only, and are
reclaimed by
-Iceberg's normal `remove_orphan_files` maintenance.
-
-A failure during the driver-side commit itself behaves exactly as on the stock
path: the
-commit messages carry genuine `SparkWrite$TaskCommit` objects, so Iceberg's own
-`SparkWrite.abort` cleanup (which deletes the files listed in the commit
messages for
-cleanable failures) applies unchanged.
+commit operator aborts without committing anything. A failed task attempt also
deletes the
+data files it created, as iceberg-java's writer abort does. The native writer
records every
+location it hands to a file writer, and exactly one side owns deleting them at
any moment: the
+native writer owns them until its output batch reaches the JVM (so it cleans
up a failed write,
+a task torn down before the write completed — for example because the operator
feeding it threw
+— and a failure encoding the manifest or building that batch), and the JVM
owns them from then
+on through a task failure listener. The handoff does not depend on decoding
the manifest: the
+native side reports the locations in the output batch alongside it, and the
listener is handed
+them before the manifest is decoded, so a failure in that decode still cleans
up. Both
+deletions are best-effort and never mask the original failure; anything they
miss is invisible
+to every reader, since readers resolve files through committed manifests only,
and is reclaimed
+by Iceberg's normal `remove_orphan_files` maintenance.
+
+When one task fails, the tasks that had already completed leave
committed-nothing data files
+too. The committer collects each task's commit message as that task finishes,
so on a job
+failure it aborts with the completed messages and deletes their data files
through the table
+`FileIO`. (Iceberg's own `SparkWrite.abort` skips cleanup unless a commit
failed with a
+cleanable error, so on the stock path those files are left for
`remove_orphan_files`.) A
+failure during the driver-side commit itself behaves exactly as on the stock
path: the commit
+messages carry genuine `SparkWrite$TaskCommit` objects, so Iceberg's own
`SparkWrite.abort`
+cleanup (which deletes the files listed in the commit messages for cleanable
failures) applies
+unchanged.
## Accepted divergences behind the toggle
diff --git a/native/core/src/execution/operators/iceberg_write.rs
b/native/core/src/execution/operators/iceberg_write.rs
index 3a988e9838..4203712831 100644
--- a/native/core/src/execution/operators/iceberg_write.rs
+++ b/native/core/src/execution/operators/iceberg_write.rs
@@ -26,7 +26,7 @@
use std::collections::HashMap;
use std::fmt;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
use arrow::array::{ArrayRef, BinaryArray, RecordBatch, UInt32Array};
use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
@@ -47,13 +47,16 @@ use futures::TryStreamExt;
use iceberg::arrow::{
arrow_struct_to_literal, PartitionValueCalculator,
RecordBatchPartitionSplitter,
};
+use iceberg::io::FileIO;
use iceberg::spec::{
DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey,
PartitionSpec,
PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef,
Struct as IcebergStruct, StructType,
};
use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
-use iceberg::writer::file_writer::location_generator::DefaultFileNameGenerator;
+use iceberg::writer::file_writer::location_generator::{
+ DefaultFileNameGenerator, LocationGenerator,
+};
use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
use iceberg::writer::file_writer::ParquetWriterBuilder;
use iceberg::writer::partitioning::clustered_writer::ClusteredWriter;
@@ -78,8 +81,146 @@ use crate::execution::operators::iceberg_partition_path::{
};
/// Builder chain instantiated once per task and handed to the partitioning
wrapper.
-type IcebergDataFileWriterBuilder =
- DataFileWriterBuilder<ParquetWriterBuilder, CometLocationGenerator,
DefaultFileNameGenerator>;
+type IcebergDataFileWriterBuilder = DataFileWriterBuilder<
+ ParquetWriterBuilder,
+ TrackingLocationGenerator,
+ DefaultFileNameGenerator,
+>;
+
+/// [`CometLocationGenerator`] 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: CometLocationGenerator,
+ locations: Arc<Mutex<Vec<String>>>,
+}
+
+impl TrackingLocationGenerator {
+ fn new(inner: CometLocationGenerator) -> Self {
+ Self {
+ inner,
+ 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.
+ ///
+ /// Ownership is given up only after the deletes have finished. Clearing
`armed` first would
+ /// remove the cancellation fallback while deletion is still in progress:
if this future is
+ /// dropped after one delete has yielded, `Drop` would see a disarmed
guard, return
+ /// immediately, and leave the remaining files with no owner. Re-deleting
a file that the
+ /// cancelled run already removed is harmless, since `delete_task_files`
is best effort and
+ /// logs rather than fails.
+ async fn abort(&mut self) {
+ delete_task_files(&self.file_io, self.generator.locations()).await;
+ self.armed = false;
+ }
+}
+
+impl Drop for AbortOnDrop {
+ fn drop(&mut self) {
+ if !self.armed {
+ return;
+ }
+ let locations = self.generator.locations();
+ if locations.is_empty() {
+ return;
+ }
+ let file_io = self.file_io.clone();
+ match tokio::runtime::Handle::try_current() {
+ // Dropped from inside a runtime (the stream was dropped while
being polled): the
+ // delete cannot block here, so hand it to the runtime.
+ Ok(handle) => {
+ handle.spawn(async move { delete_task_files(&file_io,
locations).await });
+ }
+ // Dropped from a plain JVM thread (`releasePlan`): run the delete
to completion so the
+ // files are gone before the task reports its failure.
+ Err(_) => match tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ {
+ Ok(runtime) => runtime.block_on(delete_task_files(&file_io,
locations)),
+ Err(e) => log::warn!(
+ "Could not build a runtime to delete {} data file(s) left
by a failed \
+ Iceberg write task: {e}",
+ locations.len()
+ ),
+ },
+ }
+ }
+}
+
+/// Best-effort deletion of every file a failed task attempt created, the
native counterpart of
+/// iceberg-java's `SparkCleanupUtil.deleteTaskFiles`. Failures are logged
rather than returned:
+/// the original task failure must stay the one Spark reports, and anything
left behind is still
+/// invisible to readers and reclaimed by `remove_orphan_files`.
+async fn delete_task_files(file_io: &FileIO, locations: Vec<String>) {
+ if locations.is_empty() {
+ return;
+ }
+ let mut deleted = 0usize;
+ for location in &locations {
+ match file_io.delete(location).await {
+ Ok(()) => deleted += 1,
+ Err(e) => log::warn!(
+ "Failed to delete data file {location} left by a failed
Iceberg write task: {e}"
+ ),
+ }
+ }
+ log::info!(
+ "Deleted {deleted} of {} data file(s) left by a failed Iceberg write
task",
+ locations.len()
+ );
+}
/// How many rows the rolling writer may take before it re-checks the target
file size.
///
@@ -231,7 +372,7 @@ impl ExecutionPlan for IcebergWriteExec {
let output_schema = Arc::clone(&self.output_schema);
let task = async move {
- let data_files = run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
input_stream,
Arc::clone(&common),
Arc::clone(&iceberg_schema),
@@ -243,17 +384,35 @@ impl ExecutionPlan for IcebergWriteExec {
write_time,
)
.await?;
- let manifest_bytes = encode_data_files_as_manifest(
- data_files,
- iceberg_schema,
- partition_spec,
- partition_id,
- task_attempt_id,
- &common.operation_id,
- )
- .await?;
- let batch = build_output_batch(manifest_bytes, &output_schema)?;
- Ok::<_, DataFusionError>(futures::stream::iter(vec![Ok(batch)]))
+ // The guard is still armed: until the output batch reaches the
JVM nothing else knows
+ // which files this attempt created, so a failure while packaging
it has to delete
+ // them here.
+ let locations = abort_guard.locations();
+ let packaged = async {
+ let manifest_bytes = encode_data_files_as_manifest(
+ data_files,
+ iceberg_schema,
+ partition_spec,
+ partition_id,
+ task_attempt_id,
+ &common.operation_id,
+ )
+ .await?;
+ build_output_batch(manifest_bytes, &locations, &output_schema)
+ }
+ .await;
+ match packaged {
+ // The batch carries the locations, and the JVM takes cleanup
ownership of them
+ // before it decodes the manifest, so the guard's job is done.
+ Ok(batch) => {
+ abort_guard.disarm();
+ Ok::<_,
DataFusionError>(futures::stream::iter(vec![Ok(batch)]))
+ }
+ Err(e) => {
+ abort_guard.abort().await;
+ Err(e)
+ }
+ }
};
Ok(Box::pin(RecordBatchStreamAdapter::new(
@@ -294,6 +453,9 @@ impl DisplayAs for IcebergWriteExec {
/// batch with `PARQUET_FIELD_ID_META_KEY` metadata so iceberg-rust can match
Arrow columns to
/// Iceberg field IDs, and routes through
`UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter`
/// depending on `writer_mode`.
+///
+/// On success the still-armed [`AbortOnDrop`] is returned along with the data
files: the caller
+/// owns cleanup until the output batch has been handed to the JVM.
#[allow(clippy::too_many_arguments)]
async fn run_write_task(
mut input: SendableRecordBatchStream,
@@ -305,7 +467,7 @@ async fn run_write_task(
partition_id: Option<i32>,
task_attempt_id: Option<i64>,
write_time: Time,
-) -> DFResult<Vec<DataFile>> {
+) -> DFResult<(Vec<DataFile>, AbortOnDrop)> {
// The JVM exec wrapper stamps both ids per task; a missing id means the
plan template was
// executed directly, and defaulting would make every task collide on the
same file names.
let partition_id = partition_id.ok_or_else(|| {
@@ -329,12 +491,14 @@ async fn run_write_task(
// Resolves the write's partition type once per task, before any data is
written, so a spec the
// location generator could not render fails the task cleanly rather than
panicking inside the
// infallible `LocationGenerator::generate_location`.
- let location_generator = CometLocationGenerator::try_new(
- common.data_location.clone(),
- &partition_spec,
- &iceberg_schema,
- )
- .map_err(iceberg_err)?;
+ let location_generator = TrackingLocationGenerator::new(
+ CometLocationGenerator::try_new(
+ common.data_location.clone(),
+ &partition_spec,
+ &iceberg_schema,
+ )
+ .map_err(iceberg_err)?,
+ );
let file_name_generator = DefaultFileNameGenerator::new(
file_name_prefix(partition_id, task_attempt_id, &common.operation_id),
None,
@@ -344,11 +508,16 @@ async fn run_write_task(
let rolling_builder = RollingFileWriterBuilder::new(
parquet_builder,
common.target_file_size_bytes as usize,
- file_io,
- location_generator,
+ file_io.clone(),
+ location_generator.clone(),
file_name_generator,
);
let data_file_builder = DataFileWriterBuilder::new(rolling_builder);
+ let mut abort_guard = AbortOnDrop {
+ file_io,
+ generator: location_generator,
+ armed: true,
+ };
// Build the field-id-decorated target schema once per task; every batch
is cast against it.
let target_schema =
@@ -394,20 +563,36 @@ async fn run_write_task(
_ => None,
};
- 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,
+ slicer,
+ fanout_splitter.as_ref(),
+ clustered_splitter.as_ref(),
+ )
+ .await?;
+ }
let _timer = write_time.timer();
- writer
- .write(
- decorated,
- slicer,
- fanout_splitter.as_ref(),
- clustered_splitter.as_ref(),
- )
- .await?;
+ writer.close(clustered_splitter.as_ref()).await
+ }
+ .await;
+ match outcome {
+ // Hand the still-armed guard to the caller: manifest encoding and
output-batch
+ // construction can still fail, and until the batch reaches the JVM
nothing else knows
+ // which files this attempt created.
+ Ok(data_files) => Ok((data_files, abort_guard)),
+ // 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.
+ Err(e) => {
+ abort_guard.abort().await;
+ Err(e)
+ }
}
- let _timer = write_time.timer();
- writer.close(clustered_splitter.as_ref()).await
}
/// Enum-based dispatch over the three iceberg-rust partitioning writers, each
paired with the
@@ -646,12 +831,19 @@ fn format_partition_spec(spec: &PartitionSpec) -> String {
}
}
+/// The per-task output: the V2 data manifest the JVM decodes into
`DataFile`s, plus the plain
+/// list of locations the task's writers were handed.
+///
+/// The locations are redundant with the manifest, and deliberately so: they
let the JVM take
+/// cleanup ownership of the written files with a `ByteBuffer` walk, before it
runs the far more
+/// allocation-hungry Avro decode that recovers the `DataFile`s. A failure in
that decode would
+/// otherwise leave files that only the failed decode could have named. See
`decodeLocations` and
+/// `WrittenFileCleanup` on the JVM side (`CometIcebergWriteExec`).
fn build_output_schema() -> SchemaRef {
- Arc::new(ArrowSchema::new(vec![Field::new(
- "iceberg_manifest",
- DataType::Binary,
- false,
- )]))
+ Arc::new(ArrowSchema::new(vec![
+ Field::new("iceberg_manifest", DataType::Binary, false),
+ Field::new("written_file_locations", DataType::Binary, false),
+ ]))
}
/// Align an input batch with the field-id-decorated target schema by casting
each column. The
@@ -1025,9 +1217,31 @@ fn manifest_partition_spec(partition_spec:
&PartitionSpecRef) -> PartitionSpec {
}
}
-fn build_output_batch(manifest_bytes: Vec<u8>, output_schema: &SchemaRef) ->
DFResult<RecordBatch> {
- let array: ArrayRef =
Arc::new(BinaryArray::from(vec![manifest_bytes.as_slice()]));
- RecordBatch::try_new(Arc::clone(output_schema),
vec![array]).map_err(DataFusionError::from)
+/// Frame the written-file locations for the `written_file_locations` column:
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.
+fn encode_locations(locations: &[String]) -> Vec<u8> {
+ let mut encoded = Vec::with_capacity(
+ 4 + locations.len() * 4 + locations.iter().map(|l|
l.len()).sum::<usize>(),
+ );
+ encoded.extend_from_slice(&(locations.len() as i32).to_be_bytes());
+ for location in locations {
+ encoded.extend_from_slice(&(location.len() as i32).to_be_bytes());
+ encoded.extend_from_slice(location.as_bytes());
+ }
+ encoded
+}
+
+fn build_output_batch(
+ manifest_bytes: Vec<u8>,
+ locations: &[String],
+ output_schema: &SchemaRef,
+) -> DFResult<RecordBatch> {
+ let manifest: ArrayRef =
Arc::new(BinaryArray::from(vec![manifest_bytes.as_slice()]));
+ let encoded_locations = encode_locations(locations);
+ let locations: ArrayRef =
Arc::new(BinaryArray::from(vec![encoded_locations.as_slice()]));
+ RecordBatch::try_new(Arc::clone(output_schema), vec![manifest, locations])
+ .map_err(DataFusionError::from)
}
/// Translate `IcebergParquetWriteSettings` into parquet-rs `WriterProperties`.
@@ -1101,6 +1315,109 @@ fn compression_from_proto(codec: i32, level:
Option<i32>) -> DFResult<Compressio
#[cfg(test)]
mod tests {
use super::*;
+ use iceberg::io::FileIOBuilder;
+ use iceberg::spec::{NestedField, PrimitiveType, Type};
+ use iceberg_storage_opendal::OpenDalStorageFactory;
+
+ /// Tracking generator over an unpartitioned spec, the layout every
tracking test writes to.
+ fn tracking_generator(data_location: &str) -> TrackingLocationGenerator {
+ let schema = Arc::new(
+ IcebergSchema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![NestedField::required(
+ 1,
+ "id",
+ Type::Primitive(PrimitiveType::Int),
+ )
+ .into()])
+ .build()
+ .unwrap(),
+ );
+ let spec =
PartitionSpec::builder(Arc::clone(&schema)).build().unwrap();
+ TrackingLocationGenerator::new(
+ CometLocationGenerator::try_new(data_location.to_string(), &spec,
&schema).unwrap(),
+ )
+ }
+
+ #[tokio::test]
+ async fn failed_task_deletes_every_tracked_location() {
+ let file_io =
FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build();
+ let generator = tracking_generator("memory:/warehouse/data");
+ // Two files the writer "created", plus one location that was handed
out but never
+ // written (the file open at the time of the failure): deleting it
must not fail.
+ let first = generator.generate_location(None,
"00000-00001-op-00001.parquet");
+ let second = generator.generate_location(None,
"00000-00001-op-00002.parquet");
+ let never_written = generator.generate_location(None,
"00000-00001-op-00003.parquet");
+ for location in [&first, &second] {
+ file_io
+ .new_output(location)
+ .unwrap()
+ .write(bytes::Bytes::from_static(b"parquet"))
+ .await
+ .unwrap();
+ }
+ assert_eq!(
+ generator.locations(),
+ vec![first.clone(), second.clone(), never_written.clone()]
+ );
+ assert!(file_io.exists(&first).await.unwrap());
+
+ delete_task_files(&file_io, generator.locations()).await;
+
+ assert!(!file_io.exists(&first).await.unwrap());
+ assert!(!file_io.exists(&second).await.unwrap());
+ assert!(!file_io.exists(&never_written).await.unwrap());
+ }
+
+ #[test]
+ fn dropping_an_armed_guard_outside_a_runtime_deletes_the_files() {
+ let file_io =
FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build();
+ let generator = tracking_generator("memory:/warehouse/data");
+ let runtime = tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap();
+ let armed = generator.generate_location(None, "armed.parquet");
+ let disarmed = generator.generate_location(None, "disarmed.parquet");
+ runtime.block_on(async {
+ for location in [&armed, &disarmed] {
+ file_io
+ .new_output(location)
+ .unwrap()
+ .write(bytes::Bytes::from_static(b"parquet"))
+ .await
+ .unwrap();
+ }
+ });
+
+ // Disarmed: the task completed, nothing may be deleted.
+ let mut guard = AbortOnDrop {
+ file_io: file_io.clone(),
+ generator: generator.clone(),
+ armed: true,
+ };
+ guard.disarm();
+ drop(guard);
+ assert!(runtime.block_on(file_io.exists(&armed)).unwrap());
+
+ // Armed and dropped from a thread without a runtime, as `releasePlan`
does: the delete
+ // runs to completion before `drop` returns.
+ drop(AbortOnDrop {
+ file_io: file_io.clone(),
+ generator,
+ armed: true,
+ });
+ assert!(!runtime.block_on(file_io.exists(&armed)).unwrap());
+ assert!(!runtime.block_on(file_io.exists(&disarmed)).unwrap());
+ }
+
+ #[test]
+ fn tracked_locations_follow_the_default_layout() {
+ let generator = tracking_generator("s3://bucket/table/data");
+ let location = generator.generate_location(None, "f.parquet");
+ assert_eq!(location, "s3://bucket/table/data/f.parquet");
+ assert_eq!(generator.locations(), vec![location]);
+ }
use datafusion_comet_proto::spark_operator::CompressionCodec as ProtoCodec;
fn base_settings() -> IcebergParquetWriteSettings {
@@ -1209,11 +1526,40 @@ mod tests {
}
#[test]
- fn build_output_schema_has_single_binary_column() {
+ fn build_output_schema_carries_the_manifest_and_the_written_locations() {
let schema = build_output_schema();
- assert_eq!(schema.fields().len(), 1);
+ assert_eq!(schema.fields().len(), 2);
assert_eq!(schema.field(0).name(), "iceberg_manifest");
assert_eq!(schema.field(0).data_type(), &DataType::Binary);
+ assert_eq!(schema.field(1).name(), "written_file_locations");
+ assert_eq!(schema.field(1).data_type(), &DataType::Binary);
+ }
+
+ #[test]
+ fn encoded_locations_round_trip_through_the_jvm_framing() {
+ // Mirrors `CometIcebergWriteExec.decodeLocations`; a path containing
a newline is decoded
+ // as one location rather than two.
+ fn decode(encoded: &[u8]) -> Vec<String> {
+ let mut out = Vec::new();
+ let count = i32::from_be_bytes(encoded[0..4].try_into().unwrap());
+ let mut at = 4usize;
+ for _ in 0..count {
+ let len = i32::from_be_bytes(encoded[at..at +
4].try_into().unwrap()) as usize;
+ at += 4;
+ out.push(String::from_utf8(encoded[at..at +
len].to_vec()).unwrap());
+ at += len;
+ }
+ assert_eq!(at, encoded.len());
+ out
+ }
+ let locations = vec![
+ "s3://bucket/table/data/00000-00001-op-00001.parquet".to_string(),
+
"file:/tmp/t/data/region=a%2Fb/00000-00001-op-00002.parquet".to_string(),
+ "file:/tmp/t/data/odd\nname.parquet".to_string(),
+ "file:/tmp/t/data/日本.parquet".to_string(),
+ ];
+ assert_eq!(decode(&encode_locations(&locations)), locations);
+ assert_eq!(decode(&encode_locations(&[])), Vec::<String>::new());
}
#[test]
@@ -1408,6 +1754,7 @@ mod tests {
};
use parquet::file::properties::WriterProperties;
use std::collections::HashMap;
+ use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;
@@ -1491,7 +1838,7 @@ mod tests {
writer_mode: ProtoIcebergWriterMode,
batches: Vec<RecordBatch>,
) -> DFResult<Vec<DataFile>> {
- run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
input_stream(batches),
common,
Arc::new(schema),
@@ -1502,7 +1849,11 @@ mod tests {
Some(0),
Time::default(),
)
- .await
+ .await?;
+ // These tests assert on the written files, so they stand in for
the JVM taking
+ // ownership of them.
+ abort_guard.disarm();
+ Ok(data_files)
}
fn with_target_file_size(
@@ -1568,6 +1919,56 @@ mod tests {
assert!(data_files[0].file_path().ends_with(".parquet"));
}
+ // Manifest encoding and output-batch construction run after the
writer has closed and can
+ // still fail. `run_write_task` therefore hands its caller a guard
that is still armed, so
+ // the files a successful writer produced are deleted if packaging
them fails.
+ #[tokio::test]
+ async fn
a_successful_write_returns_an_armed_guard_that_can_still_delete_its_files() {
+ let temp_dir = TempDir::new().unwrap();
+ let data_location = format!("file://{}",
temp_dir.path().display());
+ let schema = iceberg_user_schema();
+ let spec = PartitionSpec::builder(Arc::new(schema.clone()))
+ .build()
+ .unwrap();
+ let common = common(
+ data_location,
+ serde_json::to_string(&spec).unwrap(),
+ serde_json::to_string(&schema).unwrap(),
+ ProtoIcebergWriterMode::IcebergWriterUnpartitioned,
+ );
+
+ let (data_files, mut abort_guard) = run_write_task(
+ input_stream(vec![batch(&[1, 2, 3], &["us", "eu", "us"])]),
+ common,
+ Arc::new(schema),
+ Arc::new(spec),
+ ProtoIcebergWriterMode::IcebergWriterUnpartitioned,
+ WriterProperties::builder().build(),
+ Some(0),
+ Some(0),
+ Time::default(),
+ )
+ .await
+ .unwrap();
+
+ assert!(
+ abort_guard.armed,
+ "the caller owns cleanup until the JVM does"
+ );
+ let written: Vec<PathBuf> = data_files
+ .iter()
+ .map(|f|
PathBuf::from(f.file_path().trim_start_matches("file:")))
+ .collect();
+ assert!(written.iter().all(|p| p.exists()), "{written:?}");
+ assert_eq!(abort_guard.locations().len(), written.len());
+
+ // What the outer task does when `encode_data_files_as_manifest` or
+ // `build_output_batch` fails.
+ abort_guard.abort().await;
+ assert!(written.iter().all(|p| !p.exists()), "{written:?}");
+ assert!(!abort_guard.armed, "aborting also gives up ownership");
+ }
+
#[tokio::test]
async fn fanout_partitioned_write_produces_one_file_per_partition() {
let temp_dir = TempDir::new().unwrap();
@@ -2181,7 +2582,7 @@ mod tests {
let schema_arc = Arc::new(schema);
let spec_arc = Arc::new(spec);
- let data_files = run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
input_stream(vec![batch(&[10, 20], &["x", "y"])]),
Arc::clone(&common),
Arc::clone(&schema_arc),
@@ -2194,6 +2595,8 @@ mod tests {
)
.await
.unwrap();
+ let locations = abort_guard.locations();
+ abort_guard.disarm();
let manifest_bytes = encode_data_files_as_manifest(
data_files.clone(),
@@ -2206,8 +2609,18 @@ mod tests {
.await
.unwrap();
let output_schema = build_output_schema();
- let batch = build_output_batch(manifest_bytes.clone(),
&output_schema).unwrap();
+ let batch =
+ build_output_batch(manifest_bytes.clone(), &locations,
&output_schema).unwrap();
assert_eq!(batch.num_rows(), 1);
+ // The locations column names every file the manifest does, so the
JVM can clean up
+ // without decoding the manifest.
+ assert_eq!(
+ locations,
+ data_files
+ .iter()
+ .map(|f| f.file_path().to_string())
+ .collect::<Vec<_>>()
+ );
let manifest = Manifest::parse_avro(&manifest_bytes).unwrap();
let entries = manifest.entries();
@@ -2297,7 +2710,7 @@ mod tests {
let schema_arc = Arc::new(schema);
let spec_arc = Arc::new(spec);
- let data_files = run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
input_stream(vec![batch(&[1, 2], &["us", "eu"])]),
Arc::clone(&common),
Arc::clone(&schema_arc),
@@ -2310,6 +2723,8 @@ mod tests {
)
.await
.unwrap();
+ // The write succeeded, so the guard must not delete what it
produced.
+ abort_guard.disarm();
assert_eq!(data_files.len(), 1);
// No partition directory, matching iceberg-java's
`UnpartitionedDataWriter`: the file
// sits directly under the data location.
@@ -2362,7 +2777,7 @@ mod tests {
let schema_arc = Arc::new(schema);
let spec_arc = Arc::new(spec);
- let data_files = run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
input_stream(vec![batch(&[1], &["us"])]),
Arc::clone(&common),
Arc::clone(&schema_arc),
@@ -2375,6 +2790,8 @@ mod tests {
)
.await
.unwrap();
+ // The write succeeded, so the guard must not delete what it
produced.
+ abort_guard.disarm();
let manifest_bytes = encode_data_files_as_manifest(
data_files,
schema_arc,
@@ -2449,7 +2866,7 @@ mod tests {
)
.unwrap();
- let data_files = run_write_task(
+ let (data_files, mut abort_guard) = run_write_task(
Box::pin(RecordBatchStreamAdapter::new(
arrow_schema,
futures::stream::iter(vec![Ok::<_,
DataFusionError>(batch)]),
@@ -2465,6 +2882,8 @@ mod tests {
)
.await
.unwrap();
+ // The write succeeded, so the guard must not delete what it
produced.
+ abort_guard.disarm();
assert_eq!(data_files.len(), 1);
assert!(
@@ -2476,6 +2895,83 @@ mod tests {
);
}
}
+
+ /// Cancelling `abort()` part way through must leave the guard armed, so
`Drop` still owns
+ /// the files the cancelled run did not reach. Clearing `armed` before the
await made `Drop`
+ /// return immediately and orphaned the remainder.
+ ///
+ /// This uses the filesystem `FileIO` rather than the in-memory one on
purpose: the memory
+ /// backend completes every delete inside a single poll (measured: one
poll, zero pendings),
+ /// so a mid-deletion cancellation cannot be constructed against it.
+ #[tokio::test]
+ async fn cancelling_abort_keeps_the_guard_armed() {
+ use std::future::Future;
+ use std::task::Context;
+
+ let dir = tempfile::tempdir().unwrap();
+ let file_io =
FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Fs)).build();
+ let generator = tracking_generator(&format!("file://{}",
dir.path().display()));
+ let mut locations = Vec::new();
+ for i in 0..4 {
+ let location = generator.generate_location(None,
&format!("{i}.parquet"));
+ file_io
+ .new_output(&location)
+ .unwrap()
+ .write(bytes::Bytes::from_static(b"parquet"))
+ .await
+ .unwrap();
+ locations.push(location);
+ }
+
+ let mut guard = AbortOnDrop {
+ file_io: file_io.clone(),
+ generator,
+ armed: true,
+ };
+ {
+ // A no-op waker means nothing ever wakes the future, so the first
delete that yields
+ // strands it. Dropping it there is the cancellation.
+ let waker = futures::task::noop_waker();
+ let mut cx = Context::from_waker(&waker);
+ let mut fut = Box::pin(guard.abort());
+ assert!(
+ fut.as_mut().poll(&mut cx).is_pending(),
+ "expected the deletes to yield so the abort can be cancelled
mid-flight"
+ );
+ }
+
+ assert!(
+ guard.armed,
+ "a cancelled abort must not have given up ownership of the
remaining files"
+ );
+ let mut survived = Vec::new();
+ for location in &locations {
+ if file_io.exists(location).await.unwrap() {
+ survived.push(location.clone());
+ }
+ }
+ assert!(
+ !survived.is_empty(),
+ "the cancellation left nothing behind, so this test would pass
either way"
+ );
+
+ // The guard is still the owner, so its `Drop` has to finish the job.
Inside a runtime it
+ // spawns the delete, hence the bounded wait rather than an immediate
assertion.
+ drop(guard);
+ for _ in 0..200 {
+ let mut remaining = 0;
+ for location in &locations {
+ if file_io.exists(location).await.unwrap() {
+ remaining += 1;
+ }
+ }
+ if remaining == 0 {
+ return;
+ }
+ tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+ }
+ panic!("dropping the still-armed guard did not delete {survived:?}");
+ }
}
/// Pins Comet's Iceberg system-function kernels to iceberg-rust's partition
transforms.
diff --git
a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
index e1ac9993ed..95eb433d3a 100644
--- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
+++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
@@ -1736,6 +1736,66 @@ object IcebergReflection extends Logging {
}
}
+ /**
+ * Best-effort deletion of `locations` through the table's `FileIO`, for a
task that failed
+ * after iceberg-rust had already written them (the JVM-side counterpart of
iceberg-java's
+ * `SparkCleanupUtil.deleteTaskFiles`). Uses
`SupportsBulkOperations.deleteFiles` when the
+ * `FileIO` offers it and `FileIO.deleteFile(String)` per path otherwise.
Never throws: the
+ * original task failure must stay the one Spark reports. Returns the number
of locations
+ * deleted, or handed to the bulk delete. `context` identifies the task in
the log lines.
+ */
+ def deleteFilesQuietly(io: AnyRef, locations: Seq[String], context: String):
Int = {
+ import scala.jdk.CollectionConverters._
+ if (locations.isEmpty) return 0
+ val deleted = findMethod(io.getClass, "deleteFiles",
classOf[java.lang.Iterable[_]]) match {
+ case Some(bulkDelete) =>
+ try {
+ bulkDelete.invoke(io, locations.asJava)
+ locations.size
+ } catch {
+ case NonFatal(e) =>
+ logWarning(s"Bulk delete of ${locations.size} data file(s) failed
($context)", e)
+ 0
+ }
+ case None =>
+ findMethod(io.getClass, "deleteFile", classOf[String]) match {
+ case None =>
+ logWarning(
+ s"FileIO ${io.getClass.getName} has no deleteFile(String);
leaving " +
+ s"${locations.size} data file(s) for remove_orphan_files
($context)")
+ 0
+ case Some(deleteFile) =>
+ locations.count { location =>
+ try {
+ deleteFile.invoke(io, location)
+ true
+ } catch {
+ case NonFatal(e) =>
+ logWarning(s"Failed to delete data file $location
($context)", e)
+ false
+ }
+ }
+ }
+ }
+ logInfo(s"Deleted $deleted of ${locations.size} data file(s) ($context)")
+ deleted
+ }
+
+ /**
+ * The locations of the data files carried by a `SparkWrite$TaskCommit`
message (its
+ * package-private `files()`), or empty when `message` is not one. Used to
clean up after a
+ * write job that failed before any commit was attempted.
+ */
+ def taskCommitFileLocations(message: AnyRef): Seq[String] =
+ findMethodInHierarchy(message.getClass, "files") match {
+ case Some(files) =>
+ files.invoke(message) match {
+ case array: Array[_] => array.toSeq.flatMap(f =>
extractFileLocation(f))
+ case _ => Seq.empty
+ }
+ case None => Seq.empty
+ }
+
/** The table's `FileIO` (`table.io()`). Iceberg requires `FileIO` to be
`Serializable`. */
def getTableIO(table: Any): Option[AnyRef] =
findMethodInHierarchy(table.getClass, "io").map(_.invoke(table))
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala
index 25e0c5873c..00a5047625 100644
---
a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala
+++
b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.execution.{SparkPlan,
UnaryExecNode}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.types.BinaryType
import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.TaskFailureListener
import com.google.protobuf.CodedOutputStream
@@ -171,10 +172,10 @@ case class CometIcebergWriteExec(
.getOrElse(
throw new IllegalStateException(s"Native Iceberg write: no sort order
id=$sortOrderId"))
columnarRdd.mapPartitionsInternal { batches =>
- // Per-task: drain the native output (single-row Binary column carrying
one V2 data
- // manifest), decode it via Iceberg's `ManifestFiles.read` to recover
`DataFile`s, rebuild
- // each file's metrics into the ones iceberg-java's own writer would
have committed
- // (re-derived from the written parquet footer through the
version-matched
+ // Per-task: drain the native output (one row carrying one V2 data
manifest and the locations
+ // the task wrote), decode the manifest via Iceberg's
`ManifestFiles.read` to recover
+ // `DataFile`s, rebuild each file's metrics into the ones iceberg-java's
own writer would
+ // have committed (re-derived from the written parquet footer through
the version-matched
// `ParquetUtil.footerMetrics` + `MetricsConfig`, with float/double NaN
counts and bounds
// carried over from the rust writer) while re-applying the write's
`SortOrder` (which
// iceberg-rust doesn't stamp), wrap them in `SparkWrite$TaskCommit` via
reflection
@@ -182,7 +183,17 @@ case class CometIcebergWriteExec(
// `IcebergCommitExec` consumes it identically to the JVM-path output.
Empty manifests
// (zero data files written) are still committed; the surrounding
// `BatchWrite.commit(messages)` handles the no-op case correctly.
- val manifestBytes = drainAvroPayload(batches)
+ //
+ // Cleanup ownership crosses the native boundary here. The listener is
registered before the
+ // native payload is pulled and starts with nothing to delete;
`drainNativePayload` hands it
+ // the written-file locations off the payload's locations column, which
the native side fills
+ // in for exactly this purpose, before anything expensive runs. From
that point a failure in
+ // the manifest decode, the metrics rebuild, `TaskCommit` construction,
or serialization
+ // deletes the files the way iceberg-java's `DataWriter.abort()` would;
failures inside the
+ // native writer, or before its output reaches us, are cleaned up on the
native side.
+ val cleanup = new CometIcebergWriteExec.WrittenFileCleanup(tableIO)
+ Option(TaskContext.get()).foreach(_.addTaskFailureListener(cleanup))
+ val manifestBytes = drainNativePayload(batches, cleanup)
val proj = UnsafeProjection.create(schemaTypes)
val decoded =
if (manifestBytes.isEmpty) new java.util.ArrayList[AnyRef]()
@@ -223,8 +234,9 @@ case class CometIcebergWriteExec(
}
val numPartitions = childRDD.getNumPartitions
- // Native side emits a single Binary column carrying the per-task Avro
`DataFile` blob.
- val numOutputCols = 1
+ // Native side emits the per-task Avro `DataFile` blob and the locations
it wrote, one Binary
+ // column each (see `build_output_schema` in `iceberg_write.rs`).
+ val numOutputCols = 2
val capturedNativeOp = nativeOp
childRDD.mapPartitionsInternal { iter =>
@@ -273,10 +285,16 @@ case class CometIcebergWriteExec(
/**
* Drain a partition's columnar output. `iceberg_write.rs` always emits
exactly one batch with
- * one `BINARY` row carrying the per-task V2 data manifest (possibly empty
bytes when no data
- * files were written). The asserts guard against a future native-side
change.
+ * one row: a `BINARY` per-task V2 data manifest (possibly empty bytes when
no data files were
+ * written) and a `BINARY` framing of the locations the task's writers were
handed. The asserts
+ * guard against a future native-side change.
+ *
+ * `cleanup` takes ownership of those locations before the manifest bytes
are copied out of the
+ * off-heap batch, so neither that copy nor anything after it can orphan the
written files.
*/
- private def drainAvroPayload(batches: Iterator[ColumnarBatch]): Array[Byte]
= {
+ private def drainNativePayload(
+ batches: Iterator[ColumnarBatch],
+ cleanup: CometIcebergWriteExec.WrittenFileCleanup): Array[Byte] = {
require(batches.hasNext, "iceberg_write produced no output batch for this
task")
val batch = batches.next()
val payload =
@@ -284,6 +302,10 @@ case class CometIcebergWriteExec(
require(
batch.numRows() == 1,
s"iceberg_write expected exactly 1 row per task, got
${batch.numRows()}")
+ require(
+ batch.numCols() == 2,
+ s"iceberg_write expected 2 output columns per task, got
${batch.numCols()}")
+
cleanup.own(CometIcebergWriteExec.decodeLocations(batch.column(1).getBinary(0)))
batch.column(0).getBinary(0)
} finally {
batch.close()
@@ -292,3 +314,66 @@ case class CometIcebergWriteExec(
payload
}
}
+
+object CometIcebergWriteExec {
+
+ /**
+ * Decode the `written_file_locations` column written by `encode_locations`
in
+ * `iceberg_write.rs`: a big-endian `int` count, then a big-endian `int`
byte length and the
+ * UTF-8 bytes for each location. Kept free of Iceberg and of the Avro
decoder so cleanup
+ * ownership never depends on decoding the manifest that the task may have
failed to decode.
+ *
+ * Strict about consuming the whole column: every native write decodes it,
so a framing
+ * divergence between the two sides fails loudly here rather than silently
handing cleanup a
+ * truncated list of the files to delete.
+ */
+ def decodeLocations(encoded: Array[Byte]): Seq[String] = {
+ if (encoded.isEmpty) return Nil
+ val buffer = java.nio.ByteBuffer.wrap(encoded)
+ val count = buffer.getInt
+ val locations = new Array[String](count)
+ var i = 0
+ while (i < count) {
+ val bytes = new Array[Byte](buffer.getInt)
+ buffer.get(bytes)
+ locations(i) = new String(bytes, java.nio.charset.StandardCharsets.UTF_8)
+ i += 1
+ }
+ require(
+ !buffer.hasRemaining,
+ s"iceberg_write locations column has ${buffer.remaining()} trailing
byte(s) after " +
+ s"$count location(s)")
+ locations.toSeq
+ }
+
+ /**
+ * Deletes the data files a failed native Iceberg write task had already
written, the
+ * counterpart of iceberg-java's `DataWriter.abort()` /
`SparkCleanupUtil.deleteTaskFiles`.
+ *
+ * Registered before the task consumes the native output and left registered
for the rest of the
+ * task: a task that fails at any point after the write never contributes a
commit message, so
+ * every file it created is an orphan. It starts owning nothing and takes
ownership of the
+ * locations the native writer reported as soon as they are read off the
payload, which is why a
+ * failure in the manifest decode -- the step the locations would otherwise
have to be recovered
+ * from -- is covered. Deletion is best-effort and never throws, so the
failure Spark reports
+ * stays the real one.
+ *
+ * @param io
+ * the table's `FileIO`, resolved on the driver and shipped in the task
closure.
+ */
+ class WrittenFileCleanup(io: AnyRef) extends TaskFailureListener {
+ // Captured at construction (executor-side, inside the task) rather than
read from the
+ // listener argument, so the log line is identical however the cleanup is
triggered.
+ private val describeTask = Option(TaskContext.get())
+ .map(tc => s"failed Iceberg write task ${tc.partitionId()} (attempt
${tc.attemptNumber()})")
+ .getOrElse("failed Iceberg write task")
+
+ @volatile private var locations: Seq[String] = Nil
+
+ /** Take ownership of the locations the native writer reported. */
+ def own(written: Seq[String]): Unit = locations = written
+
+ override def onTaskFailure(context: TaskContext, error: Throwable): Unit =
+ IcebergReflection.deleteFilesQuietly(io, locations, describeTask)
+ }
+}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
index 5041c9fc7f..df3c38b3de 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
@@ -27,7 +27,7 @@ import org.apache.spark.sql.execution.{SparkPlan,
SQLExecution, UnaryExecNode}
import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.comet.iceberg.IcebergDriverMetricsShim
+import org.apache.comet.iceberg.{IcebergDriverMetricsShim, IcebergReflection}
/**
* Driver-side committer for Comet's split-operator Iceberg V2 write.
@@ -64,21 +64,44 @@ case class IcebergCommitExec(
}
private def collectAndCommit(): Seq[InternalRow] = {
- val messages: Array[WriterCommitMessage] =
- try {
- child.executeCollect().map { row =>
- IcebergWriteExec.deserializeMessage(row.getBinary(0))
- }
- } catch {
- case cause: Throwable =>
- // The write job failed; the BatchWrite contract still expects a
job-level abort.
- try batchWrite.abort(Array.empty[WriterCommitMessage])
- catch {
- case abortFailure: Throwable =>
- cause.addSuppressed(abortFailure)
+ // Collect each task's commit message as its task finishes, the way
Spark's own V2 write
+ // does, so that a job failure still knows which tasks completed.
`executeCollect` would
+ // only hand back the messages once every task succeeded.
+ val rdd = child.execute()
+ val messages = new Array[WriterCommitMessage](rdd.partitions.length)
+ try {
+ sparkContext.runJob(
+ rdd,
+ (iter: Iterator[InternalRow]) => iter.map(_.getBinary(0)).toArray,
+ (index: Int, payloads: Array[Array[Byte]]) => {
+ payloads.foreach { payload =>
+ messages(index) = IcebergWriteExec.deserializeMessage(payload)
}
- throw cause
- }
+ })
+ } catch {
+ case cause: Throwable =>
+ val completed = messages.filter(_ != null)
+ logError(
+ s"Iceberg write job failed; aborting with ${completed.length}
completed task " +
+ "message(s)",
+ cause)
+ // The BatchWrite contract expects a job-level abort. Iceberg's own
`SparkWrite.abort`
+ // only deletes files after a cleanable *commit* failure and skips
cleanup otherwise, so
+ // the data files of tasks that completed before the job failed would
stay behind even
+ // though nothing can reference them (no commit was attempted). Delete
them here; the
+ // failed task's own files are cleaned up by the task itself.
+ try batchWrite.abort(completed)
+ catch {
+ case abortFailure: Throwable =>
+ cause.addSuppressed(abortFailure)
+ }
+ try deleteCompletedTaskFiles(completed)
+ catch {
+ case cleanupFailure: Throwable =>
+ cause.addSuppressed(cleanupFailure)
+ }
+ throw cause
+ }
longMetric("numCommittedMessages").add(messages.length)
try {
@@ -100,6 +123,27 @@ case class IcebergCommitExec(
Nil
}
+ private def deleteCompletedTaskFiles(completed: Array[WriterCommitMessage]):
Unit = {
+ val locations = completed.toSeq.flatMap(m =>
IcebergReflection.taskCommitFileLocations(m))
+ if (locations.nonEmpty) {
+ val io = IcebergReflection
+ .getOuterSparkWrite(batchWrite)
+ .flatMap(IcebergReflection.getTableFromSparkWrite)
+ .flatMap(IcebergReflection.getTableIO)
+ io match {
+ case Some(fileIO) =>
+ IcebergReflection.deleteFilesQuietly(
+ fileIO,
+ locations,
+ s"job abort, ${completed.length} completed task(s)")
+ case None =>
+ logWarning(
+ s"Could not resolve the table FileIO; leaving ${locations.size}
data file(s) from " +
+ "completed tasks of a failed write job for remove_orphan_files")
+ }
+ }
+ }
+
private def postDriverMetrics(): Unit = {
val driverMetrics = IcebergDriverMetricsShim.reportDriverMetrics(write)
if (driverMetrics.nonEmpty) {
diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala
b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala
index 03f8fe28f6..8b2c55505c 100644
--- a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala
@@ -160,4 +160,30 @@ trait CometIcebergTestBase {
}
captured.toSeq
}
+
+ /**
+ * The executed plan of every query that fails while `action` runs, and the
failure `action`
+ * itself raised (`None` if it unexpectedly succeeded).
+ */
+ protected def captureFailedPlans(spark: SparkSession)(
+ action: => Unit): (Seq[SparkPlan], Option[Exception]) = {
+ val captured = mutable.Buffer.empty[SparkPlan]
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, durationNs:
Long): Unit = ()
+ override def onFailure(funcName: String, qe: QueryExecution, exception:
Exception): Unit =
+ captured += qe.executedPlan
+ }
+ spark.listenerManager.register(listener)
+ try {
+ val error =
+ try {
+ action
+ None
+ } catch { case e: Exception => Some(e) }
+ CometListenerBusUtils.waitUntilEmpty(spark.sparkContext)
+ (captured.toSeq, error)
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ }
}
diff --git
a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
index 0c6affcd30..7c7608d9f4 100644
--- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
@@ -27,8 +27,10 @@ import scala.collection.mutable
import scala.concurrent.{Await, Future}
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration.DurationInt
+import scala.jdk.CollectionConverters._
-import org.apache.spark.SparkConf
+import org.apache.spark.{SparkConf, Success}
+import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
import org.apache.spark.sql.CometTestBase
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.Row
@@ -40,6 +42,7 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType,
StructField, StructType}
import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus,
isSpark41Plus}
+import org.apache.comet.iceberg.IcebergReflection
private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan])
@@ -1598,6 +1601,243 @@ class CometIcebergWriteActionSuite
}
}
+ // A one-byte target file size makes the rolling writer finalize a file at
every roll point, and
+ // the writer's roll points sit on a 1000-row grid, so the source has to be
thousands of rows for
+ // the failing task to have finalized anything. A 1000-row Comet batch size
hands the writer one
+ // grid step per batch, so several files are already finalized when the UDF
throws on id 7000 --
+ // whatever the source's own partitioning is. iceberg-java's writer abort
deletes such files; the
+ // native path must too.
+ test("native acceleration: a failed task deletes the data files it already
finalized") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ val session = spark
+ import session.implicits._
+ (1 to 10000)
+ .map(i => (i, s"r$i", i.toDouble))
+ .toDF("id", "region", "amount")
+ .coalesce(1)
+ .createOrReplaceTempView("cleanup_src")
+ spark.udf.register(
+ "boom_on_seven_cleanup",
+ (id: Int) => {
+ if (id == 7000) throw new RuntimeException("boom")
+ id
+ })
+ val rollingProps = Some("'write.target-file-size-bytes'='1'")
+
+ withNativeEnabled(withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "1000") {
+ // Control: the same source and settings without the failure roll into
several files, so
+ // the failing run below really does have finalized files to clean up.
+ createTable(
+ warehouseDir,
+ "cleanup_control",
+ partitionSpec = "",
+ properties = rollingProps)
+ val controlPlans = capturePlans(spark) {
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.cleanup_control SELECT id, region,
amount FROM cleanup_src")
+ }
+ assert(
+ controlPlans.exists(p =>
+ collectWithSubqueries(p) { case w: CometIcebergWriteExec => w
}.nonEmpty),
+ "control write did not run natively")
+ val controlFiles = parquetFiles(dataDir("cleanup_control"))
+ assert(controlFiles.size >= 3, s"expected the writer to roll files,
got $controlFiles")
+
+ createTable(warehouseDir, "cleanup_target", partitionSpec = "",
properties = rollingProps)
+ coalesceInsert("cleanup_target", Seq((0, "seed", 0.0)))
+ val committed = parquetFiles(dataDir("cleanup_target"))
+ assert(committed.size == 1)
+ val before = countSnapshots("cleanup_target")
+
+ val (failedPlans, error) = captureFailedPlans(spark) {
+ spark.sql(s"INSERT INTO $catalog.$ns.cleanup_target " +
+ "SELECT boom_on_seven_cleanup(id), region, amount FROM
cleanup_src")
+ }
+ assert(
+ error.toSeq
+ .flatMap(exceptionChain)
+ .exists(t => Option(t.getMessage).exists(_.contains("boom"))),
+ s"expected the injected task failure to surface, got $error")
+ assert(
+ failedPlans.exists(p =>
+ collectWithSubqueries(p) { case w: CometIcebergWriteExec => w
}.nonEmpty),
+ s"failed write did not run
natively:\n${failedPlans.mkString("\n--\n")}")
+ assert(countSnapshots("cleanup_target") == before, "failed write must
not commit")
+ val remaining = parquetFiles(dataDir("cleanup_target"))
+ assert(
+ remaining == committed,
+ s"the failed task left data files behind: ${remaining -- committed}")
+ assertRows("cleanup_target", expectedIds = Seq(0))
+ })
+ }
+ }
+
+ // A three-task write where one task fails only after the other two have
finished: their
+ // commit messages reached the driver, so it is the committer's job abort,
not task cleanup,
+ // that has to remove their data files.
+ Seq(true, false).foreach { native =>
+ test(s"a failed write job deletes the data files of tasks that completed
(native=$native)") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ val table = s"job_abort_$native"
+ createTable(warehouseDir, table, partitionSpec = "")
+ coalesceInsert(table, Seq((0, "seed", 0.0)))
+ val committed = parquetFiles(dataDir(table))
+ val before = countSnapshots(table)
+ val session = spark
+ import session.implicits._
+ withTempPath { dir =>
+ (1 to 30)
+ .map(i => (i, s"r$i", i.toDouble))
+ .toDF("id", "region", "amount")
+ .repartition(3)
+ .write
+ .parquet(dir.getAbsolutePath)
+
spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("job_abort_src")
+ JobAbortGate.reset(othersToFinish = 2)
+ spark.udf.register(
+ "boom_after_others",
+ (id: Int) => {
+ if (id == 25) {
+ JobAbortGate.awaitOthers()
+ throw new RuntimeException("boom")
+ }
+ id
+ })
+ val listener = new SparkListener {
+ override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit =
+ if (taskEnd.reason == Success) JobAbortGate.taskFinished()
+ }
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ // One task per source file: openCostInBytes equal to
maxPartitionBytes stops the
+ // planner from packing two of these tiny files into one task.
+ val run = () =>
+ withSQLConf(
+ "spark.sql.files.maxPartitionBytes" -> "1048576",
+ "spark.sql.files.openCostInBytes" -> "1048576") {
+ spark.sql(s"INSERT INTO $catalog.$ns.$table " +
+ "SELECT boom_after_others(id), region, amount FROM
job_abort_src")
+ }
+ val (failedPlans, error) = captureFailedPlans(spark) {
+ if (native) withNativeEnabled(run()) else run()
+ }
+ assert(
+ error.toSeq
+ .flatMap(exceptionChain)
+ .exists(t => Option(t.getMessage).exists(_.contains("boom"))),
+ s"expected the injected task failure to surface, got $error")
+ val nativeWrites = failedPlans.flatMap(p =>
+ collectWithSubqueries(p) { case w: CometIcebergWriteExec => w })
+ assert(
+ nativeWrites.nonEmpty == native,
+ s"native=$native but the failed plans
were:\n${failedPlans.mkString("\n--\n")}")
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ assert(JobAbortGate.finished >= 2, "the gate must have seen two
completed tasks")
+ assert(countSnapshots(table) == before, "failed write must not
commit")
+ val remaining = parquetFiles(dataDir(table))
+ assert(
+ remaining == committed,
+ s"completed tasks left data files behind: ${remaining --
committed}")
+ assertRows(table, expectedIds = Seq(0))
+ }
+ }
+ }
+ }
+
+ test("deleteFilesQuietly removes data files through the table FileIO and
never throws") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "delete_quietly", partitionSpec = "PARTITIONED
BY (region)")
+ spark.sql(s"INSERT INTO $catalog.$ns.delete_quietly VALUES (1, 'us',
1.0), (2, 'eu', 2.0)")
+ val locations = spark
+ .sql(s"SELECT file_path FROM $catalog.$ns.delete_quietly.files")
+ .collect()
+ .map(_.getString(0))
+ .toSeq
+ assert(locations.size == 2)
+ val io = IcebergReflection
+ .getTableIO(loadIcebergTable(spark, catalog, ns, "delete_quietly"))
+ .getOrElse(fail("table.io() unavailable"))
+
+ val deleted = IcebergReflection.deleteFilesQuietly(
+ io,
+ locations :+ s"${locations.head}.missing",
+ "test")
+ assert(deleted == locations.size + 1, s"deleted=$deleted")
+ assert(parquetFiles(dataDir("delete_quietly")).isEmpty)
+ // Deleting the same paths again is a no-op rather than an error.
+ IcebergReflection.deleteFilesQuietly(io, locations, "test")
+ }
+ }
+
+ // The gap this closes: the locations the cleanup listener works from come
off the native
+ // payload's own locations column, not out of the manifest, so a failure
decoding that manifest
+ // -- the step the locations would otherwise have to be recovered from -- is
still covered.
+ test("the write cleanup listener deletes the locations the native writer
reported") {
+ assumeNativeAcceleration()
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "listener_cleanup", partitionSpec =
"PARTITIONED BY (region)")
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.listener_cleanup VALUES (1, 'us', 1.0), (2,
'eu', 2.0)")
+ val locations = spark
+ .sql(s"SELECT file_path FROM $catalog.$ns.listener_cleanup.files")
+ .collect()
+ .map(_.getString(0))
+ .toSeq
+ assert(locations.size == 2)
+ val io = IcebergReflection
+ .getTableIO(loadIcebergTable(spark, catalog, ns, "listener_cleanup"))
+ .getOrElse(fail("table.io() unavailable"))
+
+ // The listener identifies its task from the context captured at
construction, so its
+ // argument goes unused; it is registered before the native payload
arrives, and until it
+ // has been handed the locations it owns nothing and must leave the
table alone.
+ val cleanup = new CometIcebergWriteExec.WrittenFileCleanup(io)
+ cleanup.onTaskFailure(null, new RuntimeException("boom"))
+ assert(parquetFiles(dataDir("listener_cleanup")).size == 2)
+
+ cleanup.own(locations)
+ cleanup.onTaskFailure(null, new RuntimeException("boom"))
+ assert(parquetFiles(dataDir("listener_cleanup")).isEmpty)
+ }
+ }
+
+ test("the written-file locations column round-trips the native framing") {
+ // Mirrors `encode_locations` in `iceberg_write.rs`.
+ def encode(locations: Seq[String]): Array[Byte] = {
+ val bytes = new java.io.ByteArrayOutputStream()
+ val out = new java.io.DataOutputStream(bytes)
+ out.writeInt(locations.size)
+ locations.foreach { location =>
+ val utf8 = location.getBytes(java.nio.charset.StandardCharsets.UTF_8)
+ out.writeInt(utf8.length)
+ out.write(utf8)
+ }
+ bytes.toByteArray
+ }
+ // Explicit lengths rather than a separator, so a path is decoded as one
location whatever it
+ // contains: an escaped partition path, a newline, non-ASCII.
+ val locations = Seq(
+ "s3://bucket/t/data/00000-00001-op-00001.parquet",
+ "file:/tmp/t/data/region=a%2Fb/00000-00001-op-00002.parquet",
+ "file:/tmp/t/data/odd\nname.parquet",
+ "file:/tmp/t/data/日本.parquet")
+ assert(CometIcebergWriteExec.decodeLocations(encode(locations)) ==
locations)
+ assert(CometIcebergWriteExec.decodeLocations(encode(Nil)) == Nil)
+ assert(CometIcebergWriteExec.decodeLocations(Array.emptyByteArray) == Nil)
+ // A framing divergence between the two sides has to fail loudly rather
than hand cleanup a
+ // truncated list of files to delete.
+ val trailing = encode(locations) :+ 0.toByte
+ assert(
+ intercept[IllegalArgumentException](
+ CometIcebergWriteExec.decodeLocations(trailing)).getMessage
+ .contains("trailing byte"))
+ }
+
test("native acceleration: wide primitive types keep JVM-parity values and
manifest metrics") {
assumeNativeAcceleration()
withIcebergCatalog { _ =>
@@ -2041,6 +2281,34 @@ class CometIcebergWriteActionSuite
WriteSnapshot(countSnapshots(tableName) - before, plans)
}
+ /**
+ * The table's `data` directory, resolved from `Table.location()` rather
than from the warehouse
+ * conf: Spark caches the catalog instance per session, so a later test's
warehouse setting does
+ * not necessarily govern where its tables are created.
+ */
+ private def dataDir(tableName: String): File = {
+ val table = loadIcebergTable(spark, catalog, ns, tableName)
+ val location = table.getClass.getMethod("location").invoke(table).toString
+ val uri = new java.net.URI(location)
+ val root = if (uri.getScheme == null) new File(location) else new File(uri)
+ new File(root, "data")
+ }
+
+ /** Relative paths of every parquet file under `dir`, or empty when it does
not exist yet. */
+ private def parquetFiles(dir: File): Set[String] = {
+ if (!dir.exists()) return Set.empty
+ val root = dir.toPath
+ val stream = java.nio.file.Files.walk(root)
+ try {
+ stream
+ .iterator()
+ .asScala
+ .filter(p => p.toString.endsWith(".parquet"))
+ .map(p => root.relativize(p).toString)
+ .toSet
+ } finally stream.close()
+ }
+
private def countSnapshots(tableName: String): Long =
try {
spark
@@ -2262,6 +2530,34 @@ class CometIcebergWriteActionSuite
* Blocks the DELETE's write job between its scan-snapshot pin and its commit
so the test can
* inject a conflicting commit. Top-level so the UDF closure doesn't capture
the suite.
*/
+/**
+ * Lets the failing task of a multi-task write wait until the other tasks have
finished, so the
+ * driver has their commit messages when the job fails. Top-level so the UDF
closure doesn't
+ * capture the suite.
+ */
+private object JobAbortGate {
+ @volatile private var others = new CountDownLatch(0)
+ @volatile private var count = 0
+
+ def reset(othersToFinish: Int): Unit = {
+ others = new CountDownLatch(othersToFinish)
+ count = 0
+ }
+
+ def taskFinished(): Unit = synchronized {
+ count += 1
+ others.countDown()
+ }
+
+ def finished: Int = count
+
+ def awaitOthers(): Unit = {
+ if (!others.await(2, TimeUnit.MINUTES)) {
+ throw new IllegalStateException("JobAbortGate: the other tasks never
finished")
+ }
+ }
+}
+
private object ConflictGate {
@volatile private var scanStarted = new CountDownLatch(1)
@volatile private var writeReleased = new CountDownLatch(1)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
