2020package org .apache .comet
2121
2222import java .io .File
23- import java .nio .file .Files
23+ import java .nio .file .{ FileAlreadyExistsException , Files , Path }
2424import java .sql .Timestamp
2525import java .util .concurrent .{CountDownLatch , TimeUnit }
2626
@@ -29,11 +29,12 @@ import scala.concurrent.{Await, Future}
2929import scala .concurrent .ExecutionContext .Implicits .global
3030import scala .concurrent .duration .DurationInt
3131import scala .jdk .CollectionConverters ._
32+ import scala .util .control .NonFatal
3233
3334import org .json4s .{DefaultFormats , Formats }
3435import 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 }
3738import org .apache .spark .rdd .RDD
3839import org .apache .spark .scheduler .{SparkListener , SparkListenerTaskEnd }
3940import org .apache .spark .sql .CometTestBase
@@ -55,6 +56,24 @@ import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener}
5556
5657private 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+
5877private 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