[Data] Checkpoint restore hits a memory wall at ~2-4B rows; scale ID storage/filtering via the pluggable restore interface

[wingkitlee0]

Description

Target scale: ~10B-row (~10 TB and up) datasets — large enough that the
transform being checkpointed is expensive (GPU inference, heavy feature
computation) and a mid-run failure without resume support is costly, yet
well within what Ray Data pipelines process today.

A longer-form analysis backing the numbers below (measurements, benchmarks,
and the full proposal set) is in this public Google Doc (viewable without
login; commenting requires a Google account):
https://docs.google.com/document/d/1IkqTPQCBZMn4YzazldMaoNVHJrMqRiNBu0Q2poez660/edit?tab=t.0

Job-level checkpointing's restore path loads every checkpointed row ID
into a single sorted in-memory array that every filter actor consumes:

1. All committed checkpoint files are read and coalesced with
   .repartition(num_blocks=1) (checkpoint/checkpoint_filter.py:268), then
   one task converts and np.sorts the single block
   (convert_and_sort_checkpointed_ids, checkpoint_filter.py:135-166).
2. Each CheckpointFilter actor ray.gets the whole array
   (checkpoint_filter.py:407) and is scheduled with a memory reservation
   of max(1 GiB, 1.5 × array_size)
   (_internal/planner/checkpoint/plan_read_op.py:86-91).

Two independent walls follow, for a uint64 id column (8 B/row of state):

Rows checkpointed	Sorted array	Single-worker load spike (~2.5×)	Per-actor reservation (1.5×)
1B	7.5 GiB	~19 GiB	11.2 GiB (×10 actors = 112 GiB logical)
3B	22 GiB	~56 GiB	34 GiB
10B	74.5 GiB	~190 GiB	112 GiB

On common 64 GiB workers the repartition(1) load spike fails first, at
roughly 2–4B rows. The per-actor reservation becomes unplaceable soon
after — and because each actor reserves 1.5× the array, the scheduler spreads
actors across nodes, multiplying physical plasma copies cluster-wide.
String IDs fail ~100× earlier (~100–150M rows): they become object-dtype
numpy arrays, pickled into a private ~100 B/row Python-heap copy per actor —
the failure #60200 reports at 115M rows.

Concrete scale anchor: a typical embedding dataset (uuid string, int32,
384×float32 embedding) measures ~1,590 B/row in parquet, so 10B rows is
~31,000 × 500 MB files ≈ 15.7 TB, and the transform being checkpointed is
~1,400 GPU-hours of inference. Jobs of this shape exist today and cannot use
checkpointing's default restore path.

Proposed direction: expose the interface, keep the default

Rather than prescribing a new core ID structure, make the restore path
pluggable and let scale-specific structures live in user code:

1. Step 0 (PR open): [data] Make checkpoint restore pluggable via CheckpointConfig checkpoint_filter_cls / checkpoint_manager_cls #65676 adds checkpoint_manager_cls /
   checkpoint_filter_cls to CheckpointConfig. The manager returns an
   opaque (ObjectRef, size) that a matching filter interprets, so custom
   structures need no further core changes. Default behavior unchanged.
2. Step 0.5 (small follow-up): the write side has no symmetric extension
   point — CheckpointWriter.create dispatches on backend only
   (checkpoint/checkpoint_writer.py:114), and the write plan extracts only
   the id column before the writer sees data (plan_write_op.py:239-241).
   A checkpoint_writer_cls plus optionally widening the extraction to
   [id_column] + bucket_column would let plugins also control what
   checkpoints contain (pre-sorted runs, bucket-routed files).
3. Structures then implementable as plugins, without core churn: a sorted
   Arrow array (fixes the string-ID per-actor copies of [Data] Reduce the memory usage of checkpoint #60200); a
   range-chunked array (sorted ID space as range-keyed chunks — filter actors
   fetch only chunks overlapping a block's [min_id, max_id], bounding
   per-actor working set); bucketed shards keyed by source path or a
   partition column (the [Data] Refactor: Checkpoint loading from broadcast-join to bucket-join style #61509 bucket-join direction — a manager returns a
   {bucket: ObjectRef} map and filters load only the shards a block
   touches, reducing 10B rows to tens of MB of on-demand state per shard);
   an anti-join filter ([Data] refactor checkpointfilter with anti-join #60764).
4. Independent small core fixes that help the default path regardless:
   make the 1.5× reservation dtype-aware (fixed-width IDs are zero-copy
   plasma views; the per-actor heap is ~nothing), and unbiased sampling in
   _numpy_size ([Data] _numpy_size in checkpoint filter mis-sizes object arrays with non-uniform element sizes #62709).

Relationship to existing issues (not a duplicate)

• [Data] Reduce the memory usage of checkpoint #60200 reports the string-ID memory symptom; this issue quantifies the
  general walls and proposes the delivery mechanism.
• [Data] Refactor: Checkpoint loading from broadcast-join to bucket-join style #61509 (bucket-join) and [Data] refactor checkpointfilter with anti-join #60764 (anti-join) each propose one specific
  restore structure; the pluggable interface is how either can ship and be
  evaluated without rewriting the default path.
• [Data] _numpy_size in checkpoint filter mis-sizes object arrays with non-uniform element sizes #62709 is one of the small core fixes listed above.
• [data] Make checkpoint restore pluggable via CheckpointConfig checkpoint_filter_cls / checkpoint_manager_cls #65676 is the open PR implementing step 0.

Versions / Dependencies

Ray master.