Skip to content

Commit 634e37d

Browse files
authored
test: cover native Iceberg mid-write retry cleanup (#6111)
* test: cover native Iceberg mid-write retry cleanup * test: harden native Iceberg retry cleanup coverage * test: handle Spark 3.4 decimal sum without codegen * Revert "test: handle Spark 3.4 decimal sum without codegen" This reverts commit 42ea18b.
1 parent f088a54 commit 634e37d

1 file changed

Lines changed: 193 additions & 15 deletions

File tree

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

Lines changed: 193 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
package org.apache.comet
2121

2222
import java.io.File
23-
import java.nio.file.Files
23+
import java.nio.file.{FileAlreadyExistsException, Files, Path}
2424
import java.sql.Timestamp
2525
import java.util.concurrent.{CountDownLatch, TimeUnit}
2626

@@ -29,11 +29,12 @@ import scala.concurrent.{Await, Future}
2929
import scala.concurrent.ExecutionContext.Implicits.global
3030
import scala.concurrent.duration.DurationInt
3131
import scala.jdk.CollectionConverters._
32+
import scala.util.control.NonFatal
3233

3334
import org.json4s.{DefaultFormats, Formats}
3435
import org.json4s.jackson.JsonMethods.parse
3536

36-
import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success}
37+
import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext}
3738
import org.apache.spark.rdd.RDD
3839
import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
3940
import org.apache.spark.sql.CometTestBase
@@ -55,6 +56,24 @@ import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener}
5556

5657
private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan])
5758

59+
private object IcebergTestFiles {
60+
61+
/** Relative paths of every regular parquet file under `root`. */
62+
def parquetFiles(root: Path): Set[String] = {
63+
if (!Files.exists(root)) return Set.empty
64+
val stream = Files.walk(root)
65+
try {
66+
stream
67+
.iterator()
68+
.asScala
69+
.filter(path =>
70+
Files.isRegularFile(path) && path.getFileName.toString.endsWith(".parquet"))
71+
.map(path => root.relativize(path).toString)
72+
.toSet
73+
} finally stream.close()
74+
}
75+
}
76+
5877
private case class ReportedWrite(
5978
writer: String,
6079
node: String,
@@ -69,6 +88,9 @@ class CometIcebergWriteActionSuite
6988
override protected def sparkConf: SparkConf = {
7089
super.sparkConf
7190
.set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true")
91+
// local[N,M] sets task max failures to M; the retry test needs one retry, and
92+
// spark.task.maxFailures does not override this part of a local master URL.
93+
.setMaster("local[5,2]")
7294
.set(
7395
"spark.sql.extensions",
7496
"org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
@@ -1836,6 +1858,76 @@ class CometIcebergWriteActionSuite
18361858
}
18371859
}
18381860

1861+
test("native acceleration: a mid-write failure retries without orphan files") {
1862+
assumeNativeAcceleration()
1863+
withIcebergCatalog { warehouseDir =>
1864+
val session = spark
1865+
import session.implicits._
1866+
(1 to 10000)
1867+
.map(i => (i, s"r$i", i.toDouble))
1868+
.toDF("id", "region", "amount")
1869+
.coalesce(1)
1870+
.createOrReplaceTempView("retry_src")
1871+
createTable(
1872+
warehouseDir,
1873+
"retry_target",
1874+
partitionSpec = "",
1875+
properties = Some("'write.target-file-size-bytes'='1'"))
1876+
NativeWriteRetryProbe.reset()
1877+
val dataLocation = dataDir("retry_target").getAbsolutePath
1878+
spark.udf.register(
1879+
"reject_next_native_file_once",
1880+
(id: Int) => NativeWriteRetryProbe.check(id, dataLocation))
1881+
1882+
val snapshot = withNativeEnabled {
1883+
captureWrite("retry_target") {
1884+
withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "1000") {
1885+
spark.sql(s"INSERT INTO $catalog.$ns.retry_target " +
1886+
"SELECT reject_next_native_file_once(id), region, amount FROM retry_src")
1887+
}
1888+
}
1889+
}
1890+
assert(snapshot.snapshotDelta == 1L, s"expected one snapshot, got $snapshot")
1891+
assert(
1892+
NativeWriteRetryProbe.blockerCreated,
1893+
"the storage failure was not armed: " +
1894+
NativeWriteRetryProbe.failureReason.getOrElse("unknown reason"))
1895+
assert(NativeWriteRetryProbe.retrySeen, "Spark did not run a retry attempt")
1896+
assert(NativeWriteRetryProbe.blockersRemoved, "the retry did not remove every blocker")
1897+
val failedPaths = NativeWriteRetryProbe.failedPaths
1898+
assert(failedPaths.nonEmpty, "the failing attempt had not finalized a data file")
1899+
assert(
1900+
snapshot.plans.exists(p =>
1901+
collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty),
1902+
s"retry did not use the native writer: ${snapshot.plans.mkString("\n--\n")}")
1903+
1904+
val physical = parquetFiles(dataDir("retry_target"))
1905+
val root = new File(dataLocation).toPath.toAbsolutePath
1906+
val referenced = spark
1907+
.sql(s"SELECT file_path FROM $catalog.$ns.retry_target.files")
1908+
.collect()
1909+
.map { row =>
1910+
val location = row.getString(0)
1911+
val uri = new java.net.URI(location)
1912+
val file = if (uri.getScheme == null) new File(location) else new File(uri)
1913+
root.relativize(file.toPath).toString
1914+
}
1915+
.toSet
1916+
assert(referenced.nonEmpty)
1917+
assert(physical == referenced, s"orphan files: ${physical -- referenced}")
1918+
assert(
1919+
(failedPaths intersect physical).isEmpty,
1920+
s"failed attempt files survived: $failedPaths")
1921+
assert(
1922+
(failedPaths intersect referenced).isEmpty,
1923+
s"a failed attempt file was referenced by the manifest: $failedPaths")
1924+
val counts = spark
1925+
.sql(s"SELECT count(*), count(DISTINCT id) FROM $catalog.$ns.retry_target")
1926+
.head()
1927+
assert(counts.getLong(0) == 10000L && counts.getLong(1) == 10000L)
1928+
}
1929+
}
1930+
18391931
// A three-task write where one task fails only after the other two have finished: their
18401932
// commit messages reached the driver, so it is the committer's job abort, not task cleanup,
18411933
// that has to remove their data files.
@@ -2458,19 +2550,8 @@ class CometIcebergWriteActionSuite
24582550
}
24592551

24602552
/** Relative paths of every parquet file under `dir`, or empty when it does not exist yet. */
2461-
private def parquetFiles(dir: File): Set[String] = {
2462-
if (!dir.exists()) return Set.empty
2463-
val root = dir.toPath
2464-
val stream = java.nio.file.Files.walk(root)
2465-
try {
2466-
stream
2467-
.iterator()
2468-
.asScala
2469-
.filter(p => p.toString.endsWith(".parquet"))
2470-
.map(p => root.relativize(p).toString)
2471-
.toSet
2472-
} finally stream.close()
2473-
}
2553+
private def parquetFiles(dir: File): Set[String] =
2554+
IcebergTestFiles.parquetFiles(dir.toPath)
24742555

24752556
private def countSnapshots(tableName: String): Long =
24762557
try {
@@ -2868,6 +2949,103 @@ class CometIcebergWriteActionSuite
28682949

28692950
}
28702951

2952+
/** Makes the local file store reject a later data-file write in the first native attempt. */
2953+
private object NativeWriteRetryProbe {
2954+
@volatile private var firstAttemptFiles = Set.empty[String]
2955+
@volatile private var sawRetry = false
2956+
@volatile private var blockerPaths = Set.empty[String]
2957+
@volatile private var createdBlocker = false
2958+
@volatile private var firstAttemptFailure: Option[String] = None
2959+
2960+
def reset(): Unit = synchronized {
2961+
firstAttemptFiles = Set.empty
2962+
sawRetry = false
2963+
blockerPaths = Set.empty
2964+
createdBlocker = false
2965+
firstAttemptFailure = None
2966+
}
2967+
2968+
def failedPaths: Set[String] = firstAttemptFiles
2969+
2970+
def retrySeen: Boolean = sawRetry
2971+
2972+
def blockerCreated: Boolean = createdBlocker
2973+
2974+
def blockersRemoved: Boolean = blockerPaths.isEmpty
2975+
2976+
def failureReason: Option[String] = firstAttemptFailure
2977+
2978+
private def failProbe(reason: String): Nothing = {
2979+
firstAttemptFailure = Some(reason)
2980+
throw new IllegalStateException(reason)
2981+
}
2982+
2983+
private def armBlockers(dataLocation: String): Unit = {
2984+
val root = new File(dataLocation).toPath
2985+
val files = IcebergTestFiles.parquetFiles(root)
2986+
firstAttemptFiles = files
2987+
// The unpartitioned writer has at most one file open; two paths mean at least one
2988+
// earlier file has already been finalized before this input-side failure.
2989+
if (files.size < 2) {
2990+
failProbe("native writer did not finalize before injection")
2991+
}
2992+
val numbered = files.map { relative =>
2993+
val name = new File(relative).getName
2994+
val pattern = "^(.*)-(\\d{5})\\.parquet$".r
2995+
name match {
2996+
case pattern(prefix, number) => (prefix, number.toInt)
2997+
case _ => failProbe(s"unexpected native file name: $name")
2998+
}
2999+
}
3000+
val prefixes = numbered.map(_._1)
3001+
if (prefixes.size != 1) {
3002+
failProbe(s"expected one native task prefix, got $prefixes")
3003+
}
3004+
val maxFileNumber = numbered.map(_._2).max
3005+
(1 to 4).foreach { offset =>
3006+
val number = maxFileNumber + offset
3007+
val blocker = root.resolve(s"${prefixes.head}-${"%05d".format(number)}.parquet")
3008+
try {
3009+
Files.createDirectory(blocker)
3010+
blockerPaths += blocker.toString
3011+
createdBlocker = true
3012+
} catch {
3013+
// The writer can advance between the directory walk and blocker creation. Keep
3014+
// arming later file numbers rather than turning that harmless skew into the failure.
3015+
case _: FileAlreadyExistsException =>
3016+
}
3017+
}
3018+
if (!createdBlocker) {
3019+
failProbe(s"native writer advanced past all blocker candidates after $maxFileNumber")
3020+
}
3021+
// The file:// store rejects writes at these paths with EISDIR. A small range tolerates
3022+
// the native pipeline being a few files ahead of the projection's filesystem view.
3023+
// Attempt-unique names keep the blockers clear of the second writer until it removes them.
3024+
}
3025+
3026+
def check(id: Int, dataLocation: String): Int = {
3027+
if (id == 7000) {
3028+
val attempt = TaskContext.get().attemptNumber()
3029+
if (attempt == 0) {
3030+
try armBlockers(dataLocation)
3031+
catch {
3032+
case NonFatal(e) =>
3033+
if (firstAttemptFailure.isEmpty) {
3034+
val detail = Option(e.getMessage).filter(_.nonEmpty).getOrElse(e.getClass.getName)
3035+
firstAttemptFailure = Some(detail)
3036+
}
3037+
throw e
3038+
}
3039+
} else {
3040+
blockerPaths.foreach(path => Files.deleteIfExists(new File(path).toPath))
3041+
blockerPaths = Set.empty
3042+
sawRetry = true
3043+
}
3044+
}
3045+
id
3046+
}
3047+
}
3048+
28713049
/**
28723050
* Blocks the DELETE's write job between its scan-snapshot pin and its commit so the test can
28733051
* inject a conflicting commit. Top-level so the UDF closure doesn't capture the suite.

0 commit comments

Comments
 (0)