This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6153-aa2b74c6dda87aa8264aa0bea858d32819d74500 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit aae5650cdf5b0ec98c605a3a77d2e9ac66f484fe Author: Andy Grove <[email protected]> AuthorDate: Wed Sep 23 20:54:21 2026 +0000 fix: wrap Iceberg split-write failures the way Spark does when abort fails (#6153) * fix: wrap Iceberg split-write failures the way Spark does when abort fails On every supported Spark version (3.4 through 4.2), V2TableWriteExec.writeWithV2 rethrows the original job or commit failure unchanged when the abort succeeds, and only wraps it in QueryExecutionErrors.writingJobFailedError ("Writing job failed.") when batchWrite.abort itself throws, with the abort failure attached to the cause as a suppressed exception. The wrap-every-non-fatal-failure behaviour described in the issue is Spark 3.3 and earlier, which Comet does not support, so no version shim is needed. IcebergCommitExec already rethrew the raw cause, matching Spark in the common case, but also did so when the abort failed. Throw writingJobFailedError around the cause in that case. Abort, completed-task file cleanup and suppression of abort and cleanup failures onto the cause are unchanged. Add tests that run the same failing Iceberg write (a commit-time validation failure and a failed write job) with the split operator off and on and assert the thrown exception and its cause have the same types, plus tests that drive IcebergCommitExec with a BatchWrite whose abort fails. Closes #6143. * style: drop unneeded string interpolator flagged by scalafix * refactor: fold Iceberg abort wrapping into one helper and dedupe the failure tests --- .../apache/spark/sql/comet/IcebergCommitExec.scala | 33 ++-- .../comet/CometIcebergWriteActionSuite.scala | 172 ++++++++++++++++++++- 2 files changed, 191 insertions(+), 14 deletions(-) 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 df3c38b3de..04fdfd7155 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 @@ -23,6 +23,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.connector.write.{BatchWrite, Write, WriterCommitMessage} +import org.apache.spark.sql.errors.QueryExecutionErrors 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} @@ -90,17 +91,13 @@ case class IcebergCommitExec( // 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) - } + val failure = abortAfter(completed, cause) try deleteCompletedTaskFiles(completed) catch { case cleanupFailure: Throwable => cause.addSuppressed(cleanupFailure) } - throw cause + throw failure } longMetric("numCommittedMessages").add(messages.length) @@ -111,18 +108,30 @@ case class IcebergCommitExec( } catch { case cause: Throwable => logError(s"Iceberg commit failed; aborting ${messages.length} task message(s)", cause) - try batchWrite.abort(messages) - catch { - case abortFailure: Throwable => - cause.addSuppressed(abortFailure) - } - throw cause + throw abortAfter(messages, cause) } refreshCache() Nil } + /** + * Aborts the write after `cause` and returns what the failed write throws, matching Spark's + * `V2TableWriteExec.writeWithV2` on every supported Spark version: `cause` itself when the + * abort succeeds, or, when the abort also fails, a `SparkException` ("Writing job failed.") + * wrapping `cause` with the abort failure attached to it as suppressed. + */ + private def abortAfter(messages: Array[WriterCommitMessage], cause: Throwable): Throwable = + try { + batchWrite.abort(messages) + cause + } catch { + case abortFailure: Throwable => + logError("Iceberg write abort failed") + cause.addSuppressed(abortFailure) + QueryExecutionErrors.writingJobFailedError(cause) + } + private def deleteCompletedTaskFiles(completed: Array[WriterCommitMessage]): Unit = { val locations = completed.toSeq.flatMap(m => IcebergReflection.taskCommitFileLocations(m)) if (locations.nonEmpty) { diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 9c08291284..3ab8998291 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -29,14 +29,18 @@ 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, SparkException, Success} +import org.apache.spark.rdd.RDD import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog -import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} @@ -518,6 +522,96 @@ class CometIcebergWriteActionSuite } } + test("a commit-time failure surfaces as the same exception as Spark's own write") { + assume(icebergAvailable, "Iceberg not available in classpath") + withIcebergCatalog { warehouseDir => + // A serializable overwrite validated from a snapshot older than a matching append fails + // Iceberg's commit-time validation deterministically, without any concurrency. + def failCommit(table: String): Throwable = { + createTable(warehouseDir, table, partitionSpec = "") + coalesceInsert(table, Seq((1, "us-east", 10.0))) + val validateFrom = spark + .sql(s"SELECT snapshot_id FROM $catalog.$ns.$table.snapshots") + .first() + .getLong(0) + coalesceInsert(table, Seq((2, "us-west", 20.0))) + val session = spark + import session.implicits._ + val overwrite = Seq((2, "us-west", 99.0)).toDF("id", "region", "amount") + val before = countSnapshots(table) + val e = intercept[Exception] { + overwrite + .coalesce(1) + .writeTo(s"$catalog.$ns.$table") + .option("isolation-level", "serializable") + .option("validate-from-snapshot-id", validateFrom.toString) + .overwrite($"id" === 2) + } + assert(countSnapshots(table) == before, "failed commit must not create a snapshot") + assertRows(table, expectedIds = Seq(1, 2)) + e + } + + assertFailsLikeSpark("commit_fail", expectedMessage = "conflict")(failCommit) + } + } + + test("a failed write job surfaces as the same exception as Spark's own write") { + assume(icebergAvailable, "Iceberg not available in classpath") + withIcebergCatalog { warehouseDir => + spark.udf.register( + "fail_on_seven", + (id: Int) => { + if (id == 7) throw new RuntimeException("injected task failure") + id + }) + val session = spark + import session.implicits._ + // A Parquet source, not a local relation: the optimizer would otherwise evaluate the UDF + // while folding the local relation and fail the query before any job runs. + val srcDir = new File(warehouseDir, "job_fail_src") + (1 to 10) + .map(i => (i, s"r$i", i.toDouble)) + .toDF("id", "region", "amount") + .write + .parquet(srcDir.getAbsolutePath) + spark.read.parquet(srcDir.getAbsolutePath).createOrReplaceTempView("job_fail_src") + + def failJob(table: String): Throwable = { + createTable(warehouseDir, table, partitionSpec = "") + coalesceInsert(table, Seq((1, "us-east", 10.0))) + val e = intercept[Exception] { + spark + .table("job_fail_src") + .selectExpr("fail_on_seven(id) AS id", "region", "amount") + .writeTo(s"$catalog.$ns.$table") + .append() + } + assertRows(table, expectedIds = Seq(1)) + e + } + + assertFailsLikeSpark("job_fail", expectedMessage = "injected task failure")(failJob) + } + } + + test("a failed abort after a commit failure is wrapped the way Spark wraps it") { + val commitFailure = new RuntimeException("injected commit failure") + val cause = failWithFailingAbort(FailingLeafExec(failTask = false), commitFailure) + assert(cause eq commitFailure, s"expected the commit failure as the cause, got $cause") + } + + test("a failed abort after a job failure is wrapped the way Spark wraps it") { + val cause = failWithFailingAbort( + FailingLeafExec(failTask = true), + new RuntimeException("commit must not run after a failed job")) + assert( + cause.isInstanceOf[SparkException] && + exceptionChain(cause).exists(t => + Option(t.getMessage).exists(_.contains("injected job failure"))), + s"expected the job failure as the cause, got $cause") + } + test("non-Iceberg V2 write plans through Spark unchanged with the config on") { withSQLConf( "spark.sql.catalog.testcat" -> classOf[InMemoryTableCatalog].getName, @@ -2550,6 +2644,63 @@ class CometIcebergWriteActionSuite assertRows(tableName, expectedIds) } + /** + * Runs `fail` against a fresh table on Spark's own V2 write path and then on the split plan, + * and asserts both failures mention `expectedMessage` (lower case) and that the split plan + * threw the same exception type, with the same type of cause, as Spark did. + */ + private def assertFailsLikeSpark(tablePrefix: String, expectedMessage: String)( + fail: String => Throwable): Unit = { + // Spark 3.x's `withSQLConf` returns `Unit`, hence the var. + var sparkError: Throwable = null + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + sparkError = fail(s"${tablePrefix}_spark") + } + val (plans, splitError) = captureFailedPlans(spark) { + throw fail(s"${tablePrefix}_split") + } + assert( + collectIcebergWriteOps(plans)._1.nonEmpty, + "expected the failing write to run through IcebergCommitExec:\n" + + plans.mkString("\n--\n")) + Seq(sparkError, splitError.get).foreach { e => + assert( + exceptionChain(e).exists(t => + Option(t.getMessage).exists(_.toLowerCase.contains(expectedMessage))), + s"expected a failure mentioning '$expectedMessage', got $e") + } + def shape(t: Throwable): (Class[_], Option[Class[_]]) = + (t.getClass, Option(t.getCause).map(_.getClass)) + assert( + shape(splitError.get) == shape(sparkError), + s"split plan threw ${shape(splitError.get)} but Spark's write threw ${shape(sparkError)}\n" + + s"split: ${splitError.get}\nspark: $sparkError") + } + + /** + * Runs an [[IcebergCommitExec]] over `child` whose commit throws `commitFailure` and whose + * abort fails. Asserts the failure is wrapped the way Spark's `V2TableWriteExec.writeWithV2` + * (3.4 through 4.2) wraps it, in `QueryExecutionErrors.writingJobFailedError` with the abort + * failure suppressed on the original failure, and returns that original failure. + */ + private def failWithFailingAbort(child: SparkPlan, commitFailure: Throwable): Throwable = { + val abortFailure = new RuntimeException("injected abort failure") + val batchWrite = new BatchWrite { + override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = + throw new UnsupportedOperationException + override def commit(messages: Array[WriterCommitMessage]): Unit = throw commitFailure + override def abort(messages: Array[WriterCommitMessage]): Unit = throw abortFailure + } + val exec = IcebergCommitExec(batchWrite, new Write {}, () => (), child) + val e = intercept[SparkException](exec.executeCollect()) + assert(e.getMessage.contains("Writing job failed"), s"unexpected message: ${e.getMessage}") + val cause = e.getCause + assert( + cause.getSuppressed.contains(abortFailure), + s"expected the abort failure suppressed on the cause, got ${cause.getSuppressed.toSeq}") + cause + } + private def exceptionChain(t: Throwable): Seq[Throwable] = { val chain = mutable.Buffer.empty[Throwable] var current = t @@ -2594,6 +2745,23 @@ private object JobAbortGate { } } +/** + * A leaf whose single task fails when `failTask` is set, and which otherwise produces no + * partitions, so an [[IcebergCommitExec]] above it goes straight to its commit. + */ +private case class FailingLeafExec(failTask: Boolean) extends LeafExecNode { + override def output: Seq[Attribute] = Nil + + override protected def doExecute(): RDD[InternalRow] = + if (failTask) { + sparkContext + .parallelize(Seq(0), 1) + .map[InternalRow](_ => throw new RuntimeException("injected job failure")) + } else { + sparkContext.emptyRDD[InternalRow] + } +} + private object ConflictGate { @volatile private var scanStarted = new CountDownLatch(1) @volatile private var writeReleased = new CountDownLatch(1) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
