From f0e71455be3bb20decd506d65e53b5054a8a0b2e Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 22 Sep 2026 17:43:21 +0800 Subject: [PATCH 1/4] test: cover native Iceberg mid-write retry cleanup --- .../comet/CometIcebergWriteActionSuite.scala | 138 +++++++++++++++++- 1 file changed, 137 insertions(+), 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 9c082912840..b5b19c76e94 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -29,7 +29,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 @@ -54,6 +54,7 @@ class CometIcebergWriteActionSuite override protected def sparkConf: SparkConf = { super.sparkConf .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + .setMaster("local[5,2]") .set( "spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions") @@ -1709,6 +1710,72 @@ class CometIcebergWriteActionSuite } } + test("native acceleration: a mid-write failure retries without orphan files") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val session = spark + import session.implicits._ + (1 to 10000) + .map(i => (i, s"r$i", i.toDouble)) + .toDF("id", "region", "amount") + .coalesce(1) + .createOrReplaceTempView("retry_src") + createTable( + warehouseDir, + "retry_target", + partitionSpec = "", + properties = Some("'write.target-file-size-bytes'='1'")) + NativeWriteRetryProbe.reset() + val dataLocation = dataDir("retry_target").getAbsolutePath + spark.udf.register( + "reject_next_native_file_once", + (id: Int) => NativeWriteRetryProbe.check(id, dataLocation)) + + val snapshot = withNativeEnabled { + captureWrite("retry_target") { + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "1000") { + spark.sql(s"INSERT INTO $catalog.$ns.retry_target " + + "SELECT reject_next_native_file_once(id), region, amount FROM retry_src") + } + } + } + assert(snapshot.snapshotDelta == 1L, s"expected one snapshot, got $snapshot") + assert(NativeWriteRetryProbe.blockerCreated, "the storage failure was not armed") + assert(NativeWriteRetryProbe.retrySeen, "Spark did not run a retry attempt") + val failedPaths = NativeWriteRetryProbe.failedPaths + assert(failedPaths.nonEmpty, "the failing attempt had not finalized a data file") + assert( + snapshot.plans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + s"retry did not use the native writer: ${snapshot.plans.mkString("\n--\n")}") + + val physical = parquetFiles(dataDir("retry_target")) + val root = new File(dataLocation).toPath.toAbsolutePath + val referenced = spark + .sql(s"SELECT file_path FROM $catalog.$ns.retry_target.files") + .collect() + .map { row => + val location = row.getString(0) + val uri = new java.net.URI(location) + val file = if (uri.getScheme == null) new File(location) else new File(uri) + root.relativize(file.toPath).toString + } + .toSet + assert(referenced.nonEmpty) + assert(physical == referenced, s"orphan files: ${physical -- referenced}") + assert( + (failedPaths intersect physical).isEmpty, + s"failed attempt files survived: $failedPaths") + assert( + (failedPaths intersect referenced).isEmpty, + s"a failed attempt file was referenced by the manifest: $failedPaths") + val counts = spark + .sql(s"SELECT count(*), count(DISTINCT id) FROM $catalog.$ns.retry_target") + .head() + assert(counts.getLong(0) == 10000L && counts.getLong(1) == 10000L) + } + } + // 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. @@ -2562,6 +2629,75 @@ class CometIcebergWriteActionSuite } +/** Makes the local file store reject a later data-file write in the first native attempt. */ +private object NativeWriteRetryProbe { + @volatile private var firstAttemptFiles = Set.empty[String] + @volatile private var sawRetry = false + @volatile private var blockerPath: String = null + + def reset(): Unit = synchronized { + firstAttemptFiles = Set.empty + sawRetry = false + blockerPath = null + } + + def failedPaths: Set[String] = firstAttemptFiles + + def retrySeen: Boolean = sawRetry + + def blockerCreated: Boolean = blockerPath != null + + def check(id: Int, dataLocation: String): Int = { + if (id == 7000) { + val attempt = TaskContext.get().attemptNumber() + if (attempt == 0) { + val root = new File(dataLocation).toPath + val files = if (java.nio.file.Files.exists(root)) { + 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() + } else Set.empty[String] + firstAttemptFiles = files + // The unpartitioned writer has at most one file open; two paths mean at least one + // earlier file has already been finalized before this input-side failure. + if (files.size < 2) { + throw new IllegalStateException("native writer did not finalize before injection") + } + val numbered = files.map { relative => + val name = new File(relative).getName + val pattern = "^(.*)-(\\d{5})\\.parquet$".r + name match { + case pattern(prefix, number) => (prefix, number.toInt) + case _ => throw new IllegalStateException(s"unexpected native file name: $name") + } + } + val prefixes = numbered.map(_._1) + if (prefixes.size != 1) { + throw new IllegalStateException(s"expected one native task prefix, got $prefixes") + } + val next = numbered.map(_._2).max + 1 + val blocker = root.resolve(s"${prefixes.head}-${"%05d".format(next)}.parquet") + // The file:// store will reject a write at this path with EISDIR. The blocker is + // removed during the retry; attempt-unique names keep it clear of the second writer. + java.nio.file.Files.createDirectory(blocker) + blockerPath = blocker.toString + } else { + if (blockerPath != null) { + java.nio.file.Files.deleteIfExists(new File(blockerPath).toPath) + } + sawRetry = true + } + } + id + } +} + /** * Blocks the DELETE's write job between its scan-snapshot pin and its commit so the test can * inject a conflicting commit. Top-level so the UDF closure doesn't capture the suite. From 9c4c9a85a20f77701a0b5fb9e0c3df0e5d440ce5 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 11:52:32 +0800 Subject: [PATCH 2/4] test: harden native Iceberg retry cleanup coverage --- .../comet/CometIcebergWriteActionSuite.scala | 153 +++++++++++------- 1 file changed, 98 insertions(+), 55 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index b5b19c76e94..444a82d3f89 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -20,6 +20,7 @@ package org.apache.comet import java.io.File +import java.nio.file.{FileAlreadyExistsException, Files, Path} import java.sql.Timestamp import java.util.concurrent.{CountDownLatch, TimeUnit} @@ -28,6 +29,7 @@ import scala.concurrent.{Await, Future} import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal import org.apache.spark.{SparkConf, Success, TaskContext} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} @@ -46,6 +48,24 @@ import org.apache.comet.iceberg.IcebergReflection private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan]) +private object IcebergTestFiles { + + /** Relative paths of every regular parquet file under `root`. */ + def parquetFiles(root: Path): Set[String] = { + if (!Files.exists(root)) return Set.empty + val stream = Files.walk(root) + try { + stream + .iterator() + .asScala + .filter(path => + Files.isRegularFile(path) && path.getFileName.toString.endsWith(".parquet")) + .map(path => root.relativize(path).toString) + .toSet + } finally stream.close() + } +} + class CometIcebergWriteActionSuite extends CometTestBase with AdaptiveSparkPlanHelper @@ -54,6 +74,8 @@ class CometIcebergWriteActionSuite override protected def sparkConf: SparkConf = { super.sparkConf .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + // local[N,M] sets task max failures to M; the retry test needs one retry, and + // spark.task.maxFailures does not override this part of a local master URL. .setMaster("local[5,2]") .set( "spark.sql.extensions", @@ -1740,8 +1762,12 @@ class CometIcebergWriteActionSuite } } assert(snapshot.snapshotDelta == 1L, s"expected one snapshot, got $snapshot") - assert(NativeWriteRetryProbe.blockerCreated, "the storage failure was not armed") + assert( + NativeWriteRetryProbe.blockerCreated, + "the storage failure was not armed: " + + NativeWriteRetryProbe.failureReason.getOrElse("unknown reason")) assert(NativeWriteRetryProbe.retrySeen, "Spark did not run a retry attempt") + assert(NativeWriteRetryProbe.blockersRemoved, "the retry did not remove every blocker") val failedPaths = NativeWriteRetryProbe.failedPaths assert(failedPaths.nonEmpty, "the failing attempt had not finalized a data file") assert( @@ -2398,19 +2424,8 @@ class CometIcebergWriteActionSuite } /** 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() - } + private def parquetFiles(dir: File): Set[String] = + IcebergTestFiles.parquetFiles(dir.toPath) private def countSnapshots(tableName: String): Long = try { @@ -2633,64 +2648,92 @@ class CometIcebergWriteActionSuite private object NativeWriteRetryProbe { @volatile private var firstAttemptFiles = Set.empty[String] @volatile private var sawRetry = false - @volatile private var blockerPath: String = null + @volatile private var blockerPaths = Set.empty[String] + @volatile private var createdBlocker = false + @volatile private var firstAttemptFailure: Option[String] = None def reset(): Unit = synchronized { firstAttemptFiles = Set.empty sawRetry = false - blockerPath = null + blockerPaths = Set.empty + createdBlocker = false + firstAttemptFailure = None } def failedPaths: Set[String] = firstAttemptFiles def retrySeen: Boolean = sawRetry - def blockerCreated: Boolean = blockerPath != null + def blockerCreated: Boolean = createdBlocker + + def blockersRemoved: Boolean = blockerPaths.isEmpty + + def failureReason: Option[String] = firstAttemptFailure + + private def failProbe(reason: String): Nothing = { + firstAttemptFailure = Some(reason) + throw new IllegalStateException(reason) + } + + private def armBlockers(dataLocation: String): Unit = { + val root = new File(dataLocation).toPath + val files = IcebergTestFiles.parquetFiles(root) + firstAttemptFiles = files + // The unpartitioned writer has at most one file open; two paths mean at least one + // earlier file has already been finalized before this input-side failure. + if (files.size < 2) { + failProbe("native writer did not finalize before injection") + } + val numbered = files.map { relative => + val name = new File(relative).getName + val pattern = "^(.*)-(\\d{5})\\.parquet$".r + name match { + case pattern(prefix, number) => (prefix, number.toInt) + case _ => failProbe(s"unexpected native file name: $name") + } + } + val prefixes = numbered.map(_._1) + if (prefixes.size != 1) { + failProbe(s"expected one native task prefix, got $prefixes") + } + val maxFileNumber = numbered.map(_._2).max + (1 to 4).foreach { offset => + val number = maxFileNumber + offset + val blocker = root.resolve(s"${prefixes.head}-${"%05d".format(number)}.parquet") + try { + Files.createDirectory(blocker) + blockerPaths += blocker.toString + createdBlocker = true + } catch { + // The writer can advance between the directory walk and blocker creation. Keep + // arming later file numbers rather than turning that harmless skew into the failure. + case _: FileAlreadyExistsException => + } + } + if (!createdBlocker) { + failProbe(s"native writer advanced past all blocker candidates after $maxFileNumber") + } + // The file:// store rejects writes at these paths with EISDIR. A small range tolerates + // the native pipeline being a few files ahead of the projection's filesystem view. + // Attempt-unique names keep the blockers clear of the second writer until it removes them. + } def check(id: Int, dataLocation: String): Int = { if (id == 7000) { val attempt = TaskContext.get().attemptNumber() if (attempt == 0) { - val root = new File(dataLocation).toPath - val files = if (java.nio.file.Files.exists(root)) { - 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() - } else Set.empty[String] - firstAttemptFiles = files - // The unpartitioned writer has at most one file open; two paths mean at least one - // earlier file has already been finalized before this input-side failure. - if (files.size < 2) { - throw new IllegalStateException("native writer did not finalize before injection") - } - val numbered = files.map { relative => - val name = new File(relative).getName - val pattern = "^(.*)-(\\d{5})\\.parquet$".r - name match { - case pattern(prefix, number) => (prefix, number.toInt) - case _ => throw new IllegalStateException(s"unexpected native file name: $name") - } - } - val prefixes = numbered.map(_._1) - if (prefixes.size != 1) { - throw new IllegalStateException(s"expected one native task prefix, got $prefixes") + try armBlockers(dataLocation) + catch { + case NonFatal(e) => + if (firstAttemptFailure.isEmpty) { + val detail = Option(e.getMessage).filter(_.nonEmpty).getOrElse(e.getClass.getName) + firstAttemptFailure = Some(detail) + } + throw e } - val next = numbered.map(_._2).max + 1 - val blocker = root.resolve(s"${prefixes.head}-${"%05d".format(next)}.parquet") - // The file:// store will reject a write at this path with EISDIR. The blocker is - // removed during the retry; attempt-unique names keep it clear of the second writer. - java.nio.file.Files.createDirectory(blocker) - blockerPath = blocker.toString } else { - if (blockerPath != null) { - java.nio.file.Files.deleteIfExists(new File(blockerPath).toPath) - } + blockerPaths.foreach(path => Files.deleteIfExists(new File(path).toPath)) + blockerPaths = Set.empty sawRetry = true } } From 42ea18b425a15ec31fae89c342ab9a79f4dd547c Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 17:17:25 +0800 Subject: [PATCH 3/4] test: handle Spark 3.4 decimal sum without codegen --- .../comet/exec/CometAggregateSuite.scala | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 5e80bee03cc..3dc8c675e55 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -43,7 +43,7 @@ import org.apache.spark.sql.types.{ArrayType, DataTypes, StructField, StructType import org.apache.comet.CometConf import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT -import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus} import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.RegrSparkVersions import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} @@ -1963,11 +1963,17 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { Row(null)) } withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN") { + // Spark 3.4's fully interpreted projection keeps the intermediate decimal + // unbounded and recovers 0.6. Spark 3.5+ stores it in the UnsafeRow buffer and + // latches the overflow like the other no-whole-stage-codegen paths above. + val expected = + if (isSpark35Plus) Row(null) else Row(recoveredSum.bigDecimal.setScale(38)) assertDecimalSumFallsBackLikeSpark( sql("SELECT SUM(v) FROM dec_no_codegen"), reason, ansiEnabled, - Row(null)) + expected, + ansiFailureExpected = isSpark35Plus) } } } @@ -2031,15 +2037,16 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } /** - * Asserts that `df` runs the aggregate in Spark with `reason` recorded and gives Spark's - * latched result: `expected` in legacy mode, and the same failure as Spark under ANSI. + * Asserts that `df` runs the aggregate in Spark with `reason` recorded and gives `expected`, or + * raises the same failure in both engines under ANSI when `ansiFailureExpected` is true. */ private def assertDecimalSumFallsBackLikeSpark( df: DataFrame, reason: String, ansiEnabled: Boolean, - expected: Row): Unit = { - if (ansiEnabled) { + expected: Row, + ansiFailureExpected: Boolean = true): Unit = { + if (ansiEnabled && ansiFailureExpected) { val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) assert( sparkError.isDefined && cometError.isDefined, From a6ed925acbd4c36f8fd940704da45226654efc26 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 23 Sep 2026 17:23:42 +0800 Subject: [PATCH 4/4] Revert "test: handle Spark 3.4 decimal sum without codegen" This reverts commit 42ea18b425a15ec31fae89c342ab9a79f4dd547c. --- .../comet/exec/CometAggregateSuite.scala | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 3dc8c675e55..5e80bee03cc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -43,7 +43,7 @@ import org.apache.spark.sql.types.{ArrayType, DataTypes, StructField, StructType import org.apache.comet.CometConf import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT -import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus} +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.RegrSparkVersions import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} @@ -1963,17 +1963,11 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { Row(null)) } withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN") { - // Spark 3.4's fully interpreted projection keeps the intermediate decimal - // unbounded and recovers 0.6. Spark 3.5+ stores it in the UnsafeRow buffer and - // latches the overflow like the other no-whole-stage-codegen paths above. - val expected = - if (isSpark35Plus) Row(null) else Row(recoveredSum.bigDecimal.setScale(38)) assertDecimalSumFallsBackLikeSpark( sql("SELECT SUM(v) FROM dec_no_codegen"), reason, ansiEnabled, - expected, - ansiFailureExpected = isSpark35Plus) + Row(null)) } } } @@ -2037,16 +2031,15 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } /** - * Asserts that `df` runs the aggregate in Spark with `reason` recorded and gives `expected`, or - * raises the same failure in both engines under ANSI when `ansiFailureExpected` is true. + * Asserts that `df` runs the aggregate in Spark with `reason` recorded and gives Spark's + * latched result: `expected` in legacy mode, and the same failure as Spark under ANSI. */ private def assertDecimalSumFallsBackLikeSpark( df: DataFrame, reason: String, ansiEnabled: Boolean, - expected: Row, - ansiFailureExpected: Boolean = true): Unit = { - if (ansiEnabled && ansiFailureExpected) { + expected: Row): Unit = { + if (ansiEnabled) { val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) assert( sparkError.isDefined && cometError.isDefined,