[data] Make checkpoint restore pluggable via CheckpointConfig checkpoint_filter_cls / checkpoint_manager_cls

[wingkitlee0]

Why are these changes needed?

The checkpoint restore path currently hardcodes its two concrete classes — IdColumnCheckpointManager (loads checkpointed IDs) and NumpyArrayBasedCheckpointFilter (filters blocks against them) — even though both ABCs (CheckpointManager, CheckpointFilter) already define the interfaces. This PR makes the restore path pluggable via two optional CheckpointConfig fields, checkpoint_manager_cls and checkpoint_filter_cls, mirroring the existing pluggability pattern on the write side (CheckpointWriter.create dispatches on config).

Motivation: the restore path is the main scalability pressure point for large datasets (see #60200, #61509). Concretely, to checkpoint a 10B+ row partitioned Parquet dataset, the current design doesn't scale — all checkpointed IDs are coalesced and sorted into a single in-memory array (~80 GB at 10B uint64 IDs). Supporting that scale requires checkpointing that takes advantage of the dataset's partitioning: a custom manager loads checkpoint IDs as per-partition shards instead of one array, and a matching custom filter checks each block only against its shard. This PR exposes the interfaces for that customization without changing default behavior; today it requires forking or monkeypatching internals.

Changes:

• CheckpointConfig(checkpoint_filter_cls=..., checkpoint_manager_cls=...): both optional, validated to be subclasses of their respective ABCs, default None (existing behavior unchanged).
• create_checkpoint_filter_op resolves the manager class and _CheckpointFilterFn.init_checkpoint_filter resolves the filter class from the config; covers both the V1 Read and V2 ReadFiles paths since both go through create_checkpoint_filter_op.
• CheckpointFilter, NumpyArrayBasedCheckpointFilter, CheckpointManager, and IdColumnCheckpointManager are annotated @DeveloperAPI and lazily re-exported from ray.data.checkpoint (lazy to avoid the checkpoint_filter -> ray.data.context -> ray.data.checkpoint import cycle).
• The contracts are documented on the ABCs and in the CheckpointConfig docstring: the filter is constructed per actor with (checkpoint_config, checkpointed_ids_ref); the manager is constructed on the driver with (checkpoint_config=..., data_context=...) and its load_checkpoint returns (ObjectRef, int) — the ref is passed opaquely to the filter, the size feeds the per-actor memory reservation, and (None, 0) skips the filter operator entirely.

Not duplicating existing work: no open PR touches checkpoint filter pluggability (searched gh pr list for CheckpointFilter/checkpoint filter). Related-but-different open issues: #60704 asks for custom placement of the filter in the plan, #54520 for per-op config — this PR is complementary to both (it changes which filter runs, not where).

Related issue number

Related: #60200, #61509, #60704

Checks

• I've signed off every commit (DCO).
• I've run pre-commit hooks (pre-commit run --files ...) — all hooks pass (semgrep skipped: unsupported on Linux ARM64 locally).
• Tests:
  • pytest python/ray/data/tests/test_checkpoint.py run locally: 68 passed, 3 failed — the 3 (test_partial_failure_no_duplicates{,_partitioned,_row_based}) fail identically on unmodified master in this environment (ModuleNotFoundError: No module named 'test_checkpoint' from the mid-test ray.init() when pytest runs from the repo root) and all 3 pass on both master and this branch when run from python/ray/data/tests/. Not related to this change.
  • New tests (10, all passing): TestCheckpointConfig::test_{invalid,valid}_checkpoint_filter_cls and test_{invalid,valid}_checkpoint_manager_cls (config validation), test_custom_checkpoint_filter_cls (end-to-end: a no-op filter subclass keeps already-checkpointed rows, proving the custom class is instantiated in the filter actor), and test_custom_checkpoint_manager_cls (end-to-end: a manager subclass returning (None, 0) skips filtering despite checkpoint data being present, proving the custom manager drives the restore).
• AI assistance was used for this change (Claude Code); the submitting human has reviewed every changed line and run the tests locally.

Notes for reviewers:

• Validation rejects abstract classes (inspect.isabstract) in addition to the usual issubclass check. This is stricter than Ray's other pluggable-class params (ExecutionCallback, ShuffleAggregation), which are instantiated on the driver right after validation and so fail loudly on their own. Here the filter class is serialized to and instantiated inside remote filter actors, so without early validation an abstract class only fails deep in the actor pool as a confusing crash/retry loop — validating at CheckpointConfig construction surfaces the mistake immediately on the driver.
• DeveloperAPI classes aren't required by the API-discrepancy check, so no doc/source/data/api/checkpoint.rst entry is included; happy to add one if preferred.
• The dict-form DataContext.checkpoint_config path (CheckpointConfig(**value)) picks up the new kwargs automatically.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for custom checkpoint filters and managers by allowing users to override checkpoint_filter_cls and checkpoint_manager_cls in CheckpointConfig. It includes validation to ensure the provided classes are concrete subclasses of CheckpointFilter and CheckpointManager, exposes these classes via lazy imports to prevent circular dependencies, and adds comprehensive unit tests. The review feedback suggests using a relative import instead of an absolute import in __init__.py to improve robustness, and updating the CheckpointFilter docstring to clarify that the checkpointed IDs reference is passed opaquely rather than always being a sorted NumPy array.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 55b1b63. Configure here.}