test: cover native Iceberg write handoff cleanup

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

spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala

@@ -22,14 +22,15 @@ package org.apache.comet
 import java.io.File
 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}
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.duration.DurationInt
 import scala.jdk.CollectionConverters._
 
-import org.apache.spark.{SparkConf, Success}
+import org.apache.spark.{SparkConf, Success, TaskContext}
 import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.DataFrame
@@ -1709,6 +1710,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()
+      // CometTestBase starts local[5]. Spark gives that master one allowed task failure, so
+      // the job aborts without a retry and the handoff runs once.
+      assert(handoffs.map(_._1) == Vector(0), s"expected one task: $handoffs")
+      assert(handoffs.map(_._2) == Vector(0), s"expected one attempt: $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.