diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index e79fad6c249d9..e19d562c19b5b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/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 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala index 085dbcd804665..90f991479801d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/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) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ConcurrentInMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ConcurrentInMemoryRelationSuite.scala index 161f8ec647a22..3de3cd2eb02b4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ConcurrentInMemoryRelationSuite.scala +++ b/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) + } + } + } }