[SPARK-58883][CORE] Detect truncated checkpoint directory on read via persisted partition count

[anshulbaliga7]

What changes were proposed in this pull request?

Persist the original RDD's partition count to a _num_partitions file in the checkpoint directory at write time, and validate it in ReliableCheckpointRDD.getPartitions on read. A missing file (pre-existing checkpoints) or unreadable file is tolerated for backward compatibility.

Why are the changes needed?

getPartitions only checks that part-* files are contiguous from part-00000. Deleting a middle file breaks contiguity and is caught, but deleting the trailing file leaves the rest contiguous, so the check passes and the RDD is silently read back with fewer partitions and no error. This mainly affects SparkContext.checkpointFile, which Spark Streaming recovery uses to rebuild generatedRDDs, since that path has no original RDD to compare against.

Split out of SPARK-58770 (#58004), which deliberately left this out as needing its own on-disk format.

Does this PR introduce any user-facing change?

Yes. Checkpoint directories now get one extra small metadata file. Reading back a directory missing its trailing partition file(s) now throws CHECKPOINT_RDD_PARTITION_COUNT_MISMATCH instead of silently returning fewer partitions. Directories written before this change read exactly as before.

How was this patch tested?

Added three tests to CheckpointStorageSuite:

• reading a checkpoint with the trailing partition file deleted now throws
• a checkpoint missing _num_partitions (pre-existing checkpoint) still reads back correctly
• a corrupted _num_partitions file is tolerated and still reads back correctly

Was this patch authored or co-authored using generative AI tooling?

No.

[anshulbaliga7]

@LuciferYang@dongjoon-hyun Please have a look. Thanks!! (Note: currently throws _LEGACY_ERROR_TEMP_3019 since #58004 hasn't merged yet...will rebase to CHECKPOINT_RDD_PARTITION_COUNT_MISMATCH once it does.)

[anshulbaliga7]

Gentle ping @LuciferYang@dongjoon-hyun !

[LuciferYang]

The gap this closes is real: delete the trailing part-* and the rest stay contiguous, so the old check passes and the RDD comes back a partition short. I left six comments inline; none of them block.

The one I would settle before the on-disk format ships is the first. Writing _num_partitions and reading it both tolerate any failure, but acting on it is fatal and has no switch. Two of the other comments, on the swallowed write error and on the non-atomic write, cover the remaining ways the check can end up disabled without anyone noticing.

Compatibility in the other direction looks fine: the startsWith("part-") filter in getPartitions is older than _partitioner, so an older Spark reading a directory written by this version ignores the new file. The format has no version field, though, so a later change to the payload is only safe if it also changes the filename. An older reader that still deserializes the bytes to an Int will go on comparing that Int as a partition count.

[LuciferYang]

Writing and reading _num_partitions both tolerate any failure (:312, :346), but acting on it (:93) is fatal and consults no config, and the diff does not touch internal/config. The analogous local-checkpoint integrity feature in this same release ships behind spark.checkpoint.local.verifyChecksum.enabled, an .internal() flag defaulting to true.

There is no API-level way around the throw either; neither version of checkpointFile is public. What is left is deleting _num_partitions from storage, and since the error names no directory, that means deleting it from every rdd-*, which turns the check off wholesale. Until someone does that, a streaming app fails on every restart. Gating this comparison the same way would cover it.

[anshulbaliga7]

No bypass for a fatal check. Added spark.checkpoint.verifyPartitionCount.enabled, .internal(), defaults to true, following the spark.checkpoint.local.verifyChecksum.enabled precedent you pointed to.

Setting it to false loads a truncated directory as-is instead of throwing..

[LuciferYang]

The catch in writePartitionCountToCheckpointDir never uses e (:312), so a failed write leaves one WARN with no cause. The read side treats a missing _num_partitions as a directory written by an older Spark (:343), so once the file fails to appear, truncation detection stays off for that directory while rdd.checkpoint() still returns success.

The minimum is to pass e to logWarning, add the part-file count found in the directory, and say in the scaladoc that a failed write leaves the check inactive. Letting the write throw is the stronger option, but it is not free: doCheckpoint() runs after dagScheduler.runJob has returned, so throwing fails an action whose result is already computed, and cpState stays at CheckpointingInProgress.

[anshulbaliga7]

Swallowed e in the write-failure WARN. Now passed through, and the message states plainly that truncation detection is inactive for that directory going forward.

[LuciferYang]

_num_partitions is written with no atomic replace: part files go to .part-NNNNN-attempt-K and are renamed (:217, :246), while this one, like _partitioner, is a plain fs.create (:302, :272). The difference is the consequence: losing the partitioner costs a shuffle, losing this file turns off the check the PR just added. If the driver dies mid-write, or the storage truncates it, the read at :346 swallows the NonFatal into None and one WARN is all that is left. That is the same storage this PR assumes can drop a trailing file.

Routing the Int through DataOutputStream.writeInt plus a format version, written to a temp path and renamed, makes it a fixed 8 bytes whose length shows a torn write.

[anshulbaliga7]

Non-atomic write. Rewrote to DataOutputStream.writeInt behind a 1-byte format version, written to a temp path and renamed into place. A rename failure is also logged (with cause) rather than silently dropped, so every path that disables the check now leaves a trace.

[LuciferYang]

:95 passes the same id into both RDD-id placeholders of the template, so the message reads The checkpoint of RDD 4 ... The checkpoint RDD is 4 and the original RDD's id is gone, while :199 passes originalRDD.id and newRDD.id. The second message line still carries the diagnosis hint, so it is not misleading, but neither variant names the checkpoint directory.

The write path goes through the new check too: :93 and :197 test the same condition, so the new one throws first and the old one is only reachable when _num_partitions could not be written or read back. This mismatch should get its own error condition carrying the directory path and the two counts; the neighboring INVALID_CHECKPOINT_DIRECTORY (:85) already carries path.

[anshulbaliga7]

Same RDD id in both template slots, no directory in the message. Split into a new condition, CHECKPOINT_TRUNCATED_DIRECTORY, scoped to this read-only check that carries path, expected, found. CHECKPOINT_RDD_PARTITION_COUNT_MISMATCH is untouched and still only reachable from the write-path comparison.

[LuciferYang]

Delete the write at ReliableCheckpointRDD.scala:189 and only the first of the three tests goes red. Test 2 deletes _num_partitions (:795) and test 3 overwrites it (:818), neither asserting the file exists first, and fs.delete returning false is ignored while fs.create(path, true) on a missing path is silent. So those two show that reading works, not the backward compatibility and corruption tolerance their comments claim.

Test 1 does guard its precondition (:760), and so does an older test in the same file (:707). One assert(fs.exists(countFile)) in each of tests 2 and 3 covers it, and asserting the WARN in test 3 with withLogAppender is what separates "corruption tolerated" from "the check never ran".

[anshulbaliga7]

Tests 2 and 3 didn't prove what their comments claimed so added assert(fs.exists(...)) preconditions to both, and test 3 now asserts the WARN fires via withLogAppender, so it actually distinguishes "corruption tolerated" from "the check never ran."

[LuciferYang]

The two comment lines above :302 justify overwrite = false with "a second checkpoint to the same directory would fail here", but that case is not reachable. The directory is <dir>/<UUID>/rdd-<rddId>, the UUID changes on every setCheckpointDir, rddId is unique within a SparkContext, writeRDDToCheckpointDirectory has exactly one call site (ReliableRDDCheckpointData.scala:61), and RDDCheckpointData.checkpoint() gates on cpState so it runs once.

If it were reachable the conclusion would invert too: overwrite = false keeps the first value, so every later read compares a stale count against a new directory and throws a mismatch as soon as the two disagree. "Keeps this consistent with _partitioner" is all the comment needs to say.

[anshulbaliga7]

The overwrite = false reasoning is moot now since its replaced by the temp-file-then-rename write.

[anshulbaliga7]

Thanks for the thorough review @LuciferYang , have addressed all six comments above.

Can you PTAL again? Thanks!

[anshulbaliga7]

@LuciferYang gentle ping, can you please take a look? Thanks!