From 9eddaafa6051a37d4a31ad093177c356f2499ad1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 17:06:50 -0600 Subject: [PATCH 1/2] 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 --- .../user-guide/latest/iceberg-writes.md | 14 +- .../src/execution/operators/iceberg_write.rs | 251 ++++++++++++++++-- .../comet/iceberg/IcebergReflection.scala | 45 ++++ .../sql/comet/CometIcebergWriteExec.scala | 20 ++ .../comet/CometIcebergWriteActionSuite.scala | 141 ++++++++++ 5 files changed, 446 insertions(+), 25 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index db4f82c8ed8..041531fb234 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -206,12 +206,14 @@ 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. +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 deletes them when the write fails or when the task is +torn down before the write completed (for example because the operator feeding it threw), and +once the native writer has returned, a task failure listener deletes the files named in the +decoded manifest. 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. 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 diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index 57f39d1b395..240016a3073 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -25,7 +25,7 @@ //! decodes the bytes with `ManifestFiles.read(...)` to recover the `DataFile`s for commit. 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}; @@ -45,6 +45,7 @@ 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, @@ -52,7 +53,7 @@ use iceberg::spec::{ }; use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder; use iceberg::writer::file_writer::location_generator::{ - DefaultFileNameGenerator, DefaultLocationGenerator, + DefaultFileNameGenerator, DefaultLocationGenerator, LocationGenerator, }; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; use iceberg::writer::file_writer::ParquetWriterBuilder; @@ -73,8 +74,124 @@ 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; +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>>, +} + +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 { + 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 is disarmed once the task has completed. +struct AbortOnDrop { + file_io: FileIO, + generator: TrackingLocationGenerator, + armed: bool, +} + +impl AbortOnDrop { + fn disarm(&mut self) { + 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) { + 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() + ); +} /// Native Iceberg write operator. Owns the parsed Iceberg schema/spec and the parquet writer /// properties; at task execution it builds the iceberg-rust writer stack, drains the upstream @@ -300,8 +417,7 @@ async fn run_write_task( AccessMode::Write, )?; - let location_generator = - DefaultLocationGenerator::with_data_location(common.data_location.clone()); + let location_generator = TrackingLocationGenerator::new(common.data_location.clone()); let file_name_generator = DefaultFileNameGenerator::new( file_name_prefix(partition_id, task_attempt_id, &common.operation_id), None, @@ -311,11 +427,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: file_io.clone(), + generator: location_generator.clone(), + armed: true, + }; let unpartitioned = partition_spec.is_unpartitioned(); let mut writer = match (unpartitioned, writer_mode) { @@ -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(); + outcome } /// Enum-based dispatch over the three iceberg-rust partitioning writers. Each variant takes the @@ -704,6 +835,88 @@ fn compression_from_proto(codec: i32, level: Option) -> DFResult IcebergParquetWriteSettings { 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 1c1c938d4e8..2b65834b3cc 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -1654,6 +1654,51 @@ 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 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 2cafe2cf49a..35955fabc19 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 @@ -19,6 +19,8 @@ package org.apache.spark.sql.comet +import scala.jdk.CollectionConverters._ + import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow @@ -31,6 +33,7 @@ import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, Unary 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 @@ -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 { + override def onTaskFailure(context: TaskContext, error: Throwable): Unit = { + val locations = + decoded.asScala.flatMap(f => IcebergReflection.extractFileLocation(f)).toSeq + IcebergReflection.deleteFilesQuietly( + tableIO, + locations, + s"failed Iceberg write task ${context.partitionId()} " + + s"(attempt ${context.attemptNumber()})") + } + }) + } val dataFiles = IcebergReflection.rebuildDataFilesWithJavaMetrics( decoded, tableIO, diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 514123a2520..03c7daf1dd0 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -26,6 +26,7 @@ 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.{CometListenerBusUtils, SparkConf} import org.apache.spark.sql.CometTestBase @@ -38,6 +39,7 @@ import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructFi import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} +import org.apache.comet.iceberg.IcebergReflection private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan]) @@ -1172,6 +1174,99 @@ class CometIcebergWriteActionSuite } } + // A one-byte target file size makes the rolling writer finalize a file per batch, and a two-row + // Comet batch size means the task has handed several batches to the writer before the UDF + // throws on id 7. 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 10) + .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 == 7) throw new RuntimeException("boom") + id + }) + val rollingProps = Some("'write.target-file-size-bytes'='1'") + + withNativeEnabled(withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") { + // 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.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.sql(s"INSERT INTO $catalog.$ns.cleanup_target " + + "SELECT boom_on_seven_cleanup(id), region, amount FROM cleanup_src") + } + assert( + exceptionChain(error).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)) + }) + } + } + + 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") + } + } + test("native acceleration: wide primitive types keep JVM-parity values and manifest metrics") { assumeNativeAcceleration() withIcebergCatalog { _ => @@ -1603,6 +1698,52 @@ 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() + } + + /** Runs `action`, expecting it to fail; returns the failed query plans and the failure. */ + private def captureFailedPlans(action: => Unit): (Seq[SparkPlan], Throwable) = { + 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 = intercept[Exception](action) + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + (captured.toSeq, error) + } finally { + spark.listenerManager.unregister(listener) + } + } + private def countSnapshots(tableName: String): Long = try { spark From 01c9101eab180c2301455c5b50d1d68d345cd72a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 3 Sep 2026 08:27:59 -0600 Subject: [PATCH 2/2] 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 --- .../user-guide/latest/iceberg-writes.md | 13 ++- .../comet/iceberg/IcebergReflection.scala | 15 +++ .../spark/sql/comet/IcebergCommitExec.scala | 74 ++++++++++--- .../comet/CometIcebergWriteActionSuite.scala | 104 +++++++++++++++++- 4 files changed, 186 insertions(+), 20 deletions(-) diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index 041531fb234..5701528f27b 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -215,10 +215,15 @@ decoded manifest. Both deletions are best-effort and never mask the original fai 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. -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. +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/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 2b65834b3cc..11cb9c16798 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -1699,6 +1699,21 @@ object IcebergReflection extends Logging { 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/IcebergCommitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala index 5041c9fc7f2..df3c38b3de3 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/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 03c7daf1dd0..ba7df0d6cb4 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -28,7 +28,8 @@ import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ -import org.apache.spark.{CometListenerBusUtils, SparkConf} +import org.apache.spark.{CometListenerBusUtils, SparkConf, Success} +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.Row import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec} @@ -1241,6 +1242,79 @@ class CometIcebergWriteActionSuite } } + // 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 { + if (native) withNativeEnabled(run()) else run() + } + assert( + exceptionChain(error).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 => @@ -1878,6 +1952,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)