From 042d8f3d9872d184921e2ff422181f01d451bc79 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 00:48:03 +0800 Subject: [PATCH 1/3] test: cover native Iceberg write handoff cleanup --- .../sql/comet/CometIcebergWriteExec.scala | 22 ++++- .../comet/CometIcebergWriteActionSuite.scala | 81 +++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) 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 00a5047625f..a7435f68893 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 9c082912840..ed7a05ae3cd 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -22,6 +22,7 @@ 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} @@ -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() + // This suite uses local[5,2], so Spark retries the one failed task once. Both attempts + // must reach the handoff and report files for their own cleanup listener to delete. + assert(handoffs.map(_._1).distinct == Vector(0), s"expected one task: $handoffs") + assert(handoffs.map(_._2).sorted == Vector(0, 1), s"expected two 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. From 914ae86c7276ac210e6c7c4d063e0103d741036b Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 01:36:46 +0800 Subject: [PATCH 2/3] test: import TaskContext for Iceberg handoff test --- .../scala/org/apache/comet/CometIcebergWriteActionSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index ed7a05ae3cd..be481d1d074 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -30,7 +30,7 @@ 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 From 34c32f6702bcdcd97950c1f8166fea46e313f98a Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 11:18:13 +0800 Subject: [PATCH 3/3] test: expect one handoff attempt under local[5] --- .../org/apache/comet/CometIcebergWriteActionSuite.scala | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index be481d1d074..3e99f93b8bd 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -1766,10 +1766,10 @@ class CometIcebergWriteActionSuite collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), s"failed write did not run natively:\n${failedPlans.mkString("\n--\n")}") val handoffs = attempts.get() - // This suite uses local[5,2], so Spark retries the one failed task once. Both attempts - // must reach the handoff and report files for their own cleanup listener to delete. - assert(handoffs.map(_._1).distinct == Vector(0), s"expected one task: $handoffs") - assert(handoffs.map(_._2).sorted == Vector(0, 1), s"expected two attempts: $handoffs") + // 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")