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 00a5047625..a7435f6889 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 java.util.concurrent.atomic.AtomicReference + import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow @@ -305,7 +307,9 @@ case class CometIcebergWriteExec( 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))) + val locations = CometIcebergWriteExec.decodeLocations(batch.column(1).getBinary(0)) + cleanup.own(locations) + CometIcebergWriteExec.afterNativeHandoff(locations) batch.column(0).getBinary(0) } finally { batch.close() @@ -317,6 +321,22 @@ case class CometIcebergWriteExec( object CometIcebergWriteExec { + // Local-executor test hook for the boundary between owning the native payload's paths and + // decoding its manifest. The callback is absent outside a scoped test invocation. + private val handoffFailpoint = new AtomicReference[Seq[String] => Unit]() + + private[apache] def withPostNativeHandoffFailpoint[T](callback: Seq[String] => Unit)( + body: => T): T = { + val previous = handoffFailpoint.getAndSet(callback) + try body + finally handoffFailpoint.set(previous) + } + + private[comet] def afterNativeHandoff(locations: Seq[String]): Unit = { + val callback = handoffFailpoint.get() + if (callback != null) callback(locations) + } + /** * 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 diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 01a8188f29..e25375feee 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -23,6 +23,7 @@ import java.io.File import java.nio.file.Files import java.sql.Timestamp import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable import scala.concurrent.{Await, Future} @@ -33,7 +34,7 @@ import scala.jdk.CollectionConverters._ import org.json4s.{DefaultFormats, Formats} import org.json4s.jackson.JsonMethods.parse -import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success} +import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase @@ -1836,6 +1837,86 @@ class CometIcebergWriteActionSuite } } + test("native acceleration: a post-native handoff failure cleans up task files") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + createTable(warehouseDir, "handoff_target", partitionSpec = "") + coalesceInsert("handoff_target", Seq((0, "seed", 0.0))) + val before = countSnapshots("handoff_target") + val root = dataDir("handoff_target").toPath.toAbsolutePath + + def relativePath(location: String): String = { + val uri = new java.net.URI(location) + val file = if (uri.getScheme == null) new File(location) else new File(uri) + root.relativize(file.toPath.toAbsolutePath).toString + } + + def metadataFiles: Set[String] = spark + .sql(s"SELECT file_path FROM $catalog.$ns.handoff_target.files") + .collect() + .map(row => relativePath(row.getString(0))) + .toSet + + val committed = metadataFiles + assert(committed.nonEmpty, "seed write did not create a data file") + assert(parquetFiles(root.toFile) == committed) + + val session = spark + import session.implicits._ + (1 to 1000) + .map(i => (i, s"r$i", i.toDouble)) + .toDF("id", "region", "amount") + .coalesce(1) + .createOrReplaceTempView("handoff_src") + + val attempts = new AtomicReference[Vector[(Int, Int, Vector[String])]](Vector.empty) + val (failedPlans, error) = withNativeEnabled { + CometIcebergWriteExec.withPostNativeHandoffFailpoint { locations => + val tc = TaskContext.get() + attempts.getAndUpdate(_ :+ ((tc.partitionId(), tc.attemptNumber(), locations.toVector))) + throw new RuntimeException("post-native handoff injected failure") + } { + captureFailedPlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.handoff_target " + + "SELECT id, region, amount FROM handoff_src") + } + } + } + assert( + error.toSeq + .flatMap(exceptionChain) + .exists(t => + Option(t.getMessage).exists(_.contains("post-native handoff injected failure"))), + s"expected the handoff failure to reach Spark, 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")}") + val handoffs = attempts.get() + assert(handoffs.forall(_._1 == 0), s"expected only partition 0: $handoffs") + assert( + handoffs.map(_._2) == handoffs.indices.toVector, + s"expected consecutive attempts: $handoffs") + assert( + handoffs.forall(_._3.nonEmpty), + s"native payload reported no written files: $handoffs") + val failedPaths = handoffs.flatMap(_._3).map(relativePath).toSet + + assert(countSnapshots("handoff_target") == before, "failed write must not commit") + assertRows("handoff_target", expectedIds = Seq(0)) + val physical = parquetFiles(root.toFile) + val referenced = metadataFiles + assert(physical == referenced, s"orphan files: ${physical -- referenced}") + assert(referenced == committed, s"failed write changed the table files: $referenced") + assert( + (failedPaths intersect physical).isEmpty, + s"failed task files survived: $failedPaths") + assert( + (failedPaths intersect referenced).isEmpty, + s"failed task files were committed: $failedPaths") + } + } + // 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.