Skip to content

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

Open
sunchao wants to merge 1 commit into
apache:masterfrom
sunchao:dev/chao/codex/cache-materialization-correctness
Open

[SPARK-58412][SQL] Prevent incomplete and stale cache materialization statistics#57617
sunchao wants to merge 1 commit into
apache:masterfrom
sunchao:dev/chao/codex/cache-materialization-correctness

Conversation

@sunchao

@sunchao sunchao commented Jul 28, 2026

Copy link
Copy Markdown
Member

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.

@dongjoon-hyun dongjoon-hyun left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1, LGTM (Pending CIs). Thank you, @sunchao !

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants