sunchao opened a new pull request, #57617: URL: https://github.com/apache/spark/pull/57617
### Why are the changes needed? SPARK-58412 tracks two lifecycle gaps in the cache-materialization bookkeeping introduced by SPARK-57547. That earlier change replaced raw cache-task completion counters with partition-keyed statistics so duplicate computations cannot make an `InMemoryRelation` appear materialized too early. First, `CachedRDDBuilder` currently publishes a task's partition statistics whenever the task completes successfully, even if a downstream consumer stops before exhausting the cache-building iterator. This can happen when a memory-only block cannot be stored and Spark returns the partially unrolled iterator to its consumer. Because the statistics accumulator is last-write-wins per partition, a partially consumed recomputation can replace that partition's complete row and byte counts while the cache still appears fully materialized. For example, a completed `(10 rows, N bytes)` entry can be replaced with `(0 rows, 0 bytes)` without removing the partition key, allowing AQE to treat a non-empty cache as empty. Second, `clearCache` resets the existing accumulator in place. Tasks from the retired cache generation have already captured that same accumulator, so late completions can write stale partition keys and values into the rebuilt generation. If those stale keys cover every partition, the new cache can appear complete before its own partitions finish. These are general `InMemoryRelation` correctness issues. They were identified while working on runtime-filter support in #57443, but they affect cache materialization independently of that optimizer feature. ### What changes were proposed in this PR? This change makes cache statistics represent only complete work from the current cache generation. A task now publishes its partition statistics only after the wrapped cache iterator has been exhausted and the task finishes without failure or interruption. A successful consumer that stops early therefore cannot overwrite a complete partition's statistics with partial values. Clearing a cache now installs a newly registered `PartitionKeyedAccumulator` instead of resetting the previous accumulator. Each cache-building RDD captures the accumulator for its own generation, so tasks finishing after `clearCache` can update only the retired generation. The cache contents and storage semantics are unchanged; this change only tightens the lifecycle of materialization metadata. ### How was this PR tested? The existing `CachedTableSuite` clear-cache regression now injects late updates through the retired generation's accumulator before and after rebuilding the cache. It verifies that the new generation remains unloaded until its own partitions complete and that its row and byte statistics remain exact. A new `ConcurrentInMemoryRelationSuite` regression forces a cached partition to be recomputed, then runs a successful consumer that does not consume the returned cache iterator. It verifies that the partial attempt cannot replace the complete partition statistics. The following validation passed: ``` build/sbt 'sql/testOnly org.apache.spark.sql.CachedTableSuite org.apache.spark.sql.execution.columnar.ConcurrentInMemoryRelationSuite' build/sbt sql/scalastyle sql/test:scalastyle git diff --check ``` The two affected suites ran 108 tests successfully. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
