[SPARK-58412][SQL] Prevent incomplete and stale cache materialization statistics

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala

@@ -271,12 +271,16 @@ case class CachedRDDBuilder(
   // partition has been computed -- which otherwise let AQE read rowCount 0 on a non-empty cache and
   // propagate an empty relation, silently dropping rows -- and also yields exact, de-duplicated row
   // count / size.
-  private val partitionStats: PartitionKeyedAccumulator[(Long, Long)] = {
+  private def newPartitionStats(): PartitionKeyedAccumulator[(Long, Long)] = {
     val acc = new PartitionKeyedAccumulator[(Long, Long)]
     cachedPlan.session.sparkContext.register(acc)
     acc
   }
 
+  // Old tasks can still finish after unpersist; isolating each buffer generation prevents their
+  // late updates from making a rebuilt cache appear complete.
+  private var partitionStats = newPartitionStats()
+
   val cachedName = tableName.map(n => s"In-memory table $n")
     .getOrElse(Utils.abbreviate(cachedPlan.toString, 1024))
 
@@ -296,11 +300,11 @@ case class CachedRDDBuilder(
     if (_cachedColumnBuffers != null) {
       _cachedColumnBuffers.unpersist(blocking)
       _cachedColumnBuffers = null
-      // The buffers no longer back a live RDD. Reset the one-way "loaded" latch and the keyed
-      // bookkeeping so a rebuild on this builder does not inherit a stale "loaded" state or stale
-      // statistics. Safe to reset in place: every read of the accumulator is under this monitor.
+      // The buffers no longer back a live RDD. Reset the one-way "loaded" latch and install new
+      // bookkeeping so a rebuild on this builder does not inherit stale state or late updates from
+      // tasks that captured the previous generation's accumulator.
       _cachedColumnBuffersAreLoaded = false
-      partitionStats.reset()
+      partitionStats = newPartitionStats()
     }
   }
 
@@ -338,7 +342,9 @@ case class CachedRDDBuilder(
 
   // The id of the accumulator backing this cache's materialization bookkeeping. Exposed only so
   // `CachedTableSuite`'s accumulator-cleanup test can verify it is cleared after uncache + GC.
-  private[sql] def materializationAccumulatorId: Long = partitionStats.id
+  private[sql] def materializationAccumulatorId: Long = synchronized {
+    partitionStats.id
+  }
 
   private def buildBuffers(): RDD[CachedBatch] = {
     val cb = try {
@@ -376,13 +382,18 @@ case class CachedRDDBuilder(
       // no batches).
       var localRows = 0L
       var localBytes = 0L
+      var fullyConsumed = false
       taskContext.addTaskCompletionListener[Unit] { context =>
-        if (!context.isFailed() && !context.isInterrupted()) {
+        if (fullyConsumed && !context.isFailed() && !context.isInterrupted()) {
           accumulator.add((partitionId, (localRows, localBytes)))
         }
       }
       new Iterator[CachedBatch] {
-        override def hasNext: Boolean = it.hasNext
+        override def hasNext: Boolean = {
+          val more = it.hasNext
+          if (!more) fullyConsumed = true
+          more
+        }
         override def next(): CachedBatch = {
           val batch = it.next()
           localBytes += batch.sizeInBytes

sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala

@@ -55,6 +55,7 @@ import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types.{StringType, StructField, StructType}
+import org.apache.spark.sql.util.PartitionKeyedAccumulator
 import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 import org.apache.spark.storage.StorageLevel.{MEMORY_AND_DISK_2, MEMORY_ONLY}
 import org.apache.spark.tags.SlowSQLTest
@@ -509,7 +510,7 @@ class CachedTableSuite extends SharedSparkSession
     }
   }
 
-  test("SPARK-57547: clearCache resets materialization bookkeeping") {
+  test("SPARK-57547: clearCache isolates materialization bookkeeping by cache generation") {
     val df = spark.range(0, 100, 1, numPartitions = 4).filter($"id" >= 0)
     df.cache()
     try {
@@ -523,18 +524,33 @@ class CachedTableSuite extends SharedSparkSession
       builder.cachedColumnBuffers.count()
       assert(builder.isCachedColumnBuffersLoaded)
       assert(builder.materializedRowCount == 100L)
+      val oldAccumulatorId = builder.materializationAccumulatorId
+      val oldAccumulator = AccumulatorContext.get(oldAccumulatorId).get
+        .asInstanceOf[PartitionKeyedAccumulator[(Long, Long)]]
 
       builder.clearCache()
-      // The loaded latch and the materialization stats must not survive clearCache, otherwise a
-      // rebuilt cache would inherit a stale "loaded" state with stale/zero statistics.
+      assert(builder.materializationAccumulatorId != oldAccumulatorId)
       assert(!builder.isCachedColumnBuffersLoaded)
       assert(builder.materializedRowCount == 0L)
       assert(builder.materializedSizeInBytes == 0L)
 
-      // Rebuilding works and reports correct stats again.
-      builder.cachedColumnBuffers.count()
+      // Simulate tasks from the previous cache generation completing after clearCache. Their
+      // updates must remain isolated from the new generation.
+      (0 until 4).foreach(partitionId => oldAccumulator.add((partitionId, (1L, 1L))))
+      assert(builder.materializedRowCount == 0L)
+      assert(builder.materializedSizeInBytes == 0L)
+      val rebuiltBuffers = builder.cachedColumnBuffers
+      assert(!builder.isCachedColumnBuffersLoaded)
+
+      rebuiltBuffers.count()
       assert(builder.isCachedColumnBuffersLoaded)
       assert(builder.materializedRowCount == 100L)
+      val rebuiltSizeInBytes = builder.materializedSizeInBytes
+      assert(rebuiltSizeInBytes > 0L)
+
+      oldAccumulator.add((0, (999L, 999L)))
+      assert(builder.materializedRowCount == 100L)
+      assert(builder.materializedSizeInBytes == rebuiltSizeInBytes)
     } finally {
       df.unpersist(blocking = true)
     }

sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ConcurrentInMemoryRelationSuite.scala

@@ -29,6 +29,7 @@ import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkFunSui
 import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
 import org.apache.spark.sql.columnar.CachedBatch
 import org.apache.spark.sql.functions.{lit, when}
+import org.apache.spark.storage.StorageLevel
 import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
@@ -63,11 +64,14 @@ class ConcurrentInMemoryRelationSuite extends SparkFunSuite with LocalSparkConte
     relations.head.cacheBuilder
   }
 
-  private def withSession(numExecutors: Int = 4)(
+  private def withSession(
+      numExecutors: Int = 4,
+      extraConfs: Map[String, String] = Map.empty)(
       f: SparkSession => Unit): Unit = {
     val conf = new SparkConf()
       .setMaster(s"local-cluster[$numExecutors,1,1024]")
       .setAppName("ConcurrentInMemoryRelationSuite")
+      .setAll(extraConfs)
     sc = new SparkContext(conf)
     try {
       // Wait for all executors to register so tasks spread one-per-executor as the tests assume.
@@ -322,4 +326,29 @@ class ConcurrentInMemoryRelationSuite extends SparkFunSuite with LocalSparkConte
         s"partition-keyed accumulator should report exact row count $expected, got $rowCount")
     }
   }
+
+  test("SPARK-57547: a partial cache attempt cannot overwrite complete partition statistics") {
+    withSession(
+        numExecutors = 1,
+        extraConfs = Map("spark.storage.unrollMemoryThreshold" -> (1L << 40).toString)) { spark =>
+      val cached = spark.range(0, 10, 1, numPartitions = 1).persist(StorageLevel.MEMORY_ONLY)
+      try {
+        val builder = cacheBuilderOf(cached)
+        val buffers = builder.cachedColumnBuffers
+        buffers.count()
+        assert(builder.isCachedColumnBuffersLoaded)
+        val expectedStats = (builder.materializedRowCount, builder.materializedSizeInBytes)
+        assert(expectedStats._1 == 10L)
+        assert(expectedStats._2 > 0L)
+
+        // The successful task does not consume its cache iterator. It must not replace the
+        // completed partition's statistics with (0, 0).
+        buffers.mapPartitions(_ => Iterator.empty).count()
+        assert(builder.isCachedColumnBuffersLoaded)
+        assert((builder.materializedRowCount, builder.materializedSizeInBytes) == expectedStats)
+      } finally {
+        cached.unpersist(blocking = true)
+      }
+    }
+  }
 }