diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index db4f82c8ed..041531fb23 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 57f39d1b39..240016a307 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 1c1c938d4e..2b65834b3c 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 2cafe2cf49..35955fabc1 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 514123a252..03c7daf1dd 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