[Data] Add checkpoint support for Iceberg read/write by WinkerDu · Pull Request #61753 · ray-project/ray

WinkerDu · 2026-03-15T16:36:23Z

Description

This PR adds checkpoint integration for Ray Data Iceberg read/write to improve fault tolerance during distributed execution.

On the write path, we persist row-level checkpoint IDs (via CheckpointConfig.id_column) and also checkpoint IcebergWriteResult metadata per write task. If the job fails after data files are written but before the driver commit completes, a retry can load all previously checkpointed write results and perform a single unified Iceberg commit, avoiding duplicate commits and ensuring consistency.

On the read path, we integrate the existing checkpoint filter so read tasks can skip already-processed rows by id_column. For Iceberg reads, we also support filtering planned scan files via checkpoint_path_partition_filter to reduce unnecessary work during restore.

Related issues

Fixes #59870

Additional information

Main changes:
- Write planning adds a checkpoint-writing transform that:
  - validates id_column exists in blocks
  - writes {uuid}.parquet with checkpointed IDs
  - writes {uuid}.meta.pkl containing the write task’s IcebergWriteResult for recovery
- IcebergDatasink.on_write_complete() loads checkpointed IcebergWriteResult entries (when enabled), merges them with current results, and commits once.
- IcebergDatasource.get_read_tasks() optionally filters planned files when checkpoint_path_partition_filter is provided.
Tests:
- Adds/updates python/ray/data/tests/test_checkpoint_for_iceberg.py to cover:
  - checkpoint ID + write-result metadata persistence and loading
  - failure-before-commit recovery with unified commit
  - UPSERT recovery including upsert_keys handling

Change-Id: Iee31177fe923a779b958ad18aeea738adbf273f7

Change-Id: Ib43c5ef4160fba7b567e8b970edad09568f9aa48

Change-Id: Ib148e684b6e193a307b0a3ddd02d5627ac17209a

Change-Id: I058a20dad17a1d1653befc7d52ecf29709a53e58

Change-Id: Iab693a43fbe0ff3e07f42a4039c6e76b232a1e32

Change-Id: Ie91b8235fd9d61b1521e204b66241d395b9a47ce

Change-Id: Ib0fed415eb864e140c6aaee82d584d945f0c7402

gemini-code-assist

Code Review

This pull request introduces checkpointing support for Iceberg read and write operations in Ray Data, which is a great enhancement for fault tolerance. The implementation for the write path correctly checkpoints IcebergWriteResult metadata, allowing for recovery and a unified commit on retry. The read path is also updated to filter data based on checkpoints. The changes are well-structured and include comprehensive tests for recovery scenarios. I have a few minor suggestions to improve code clarity and remove dead code.

WinkerDu · 2026-03-19T12:50:46Z

@owenowenisme @xinyuangui2 please take a review, thx :)

owenowenisme

Hey @WinkerDu, thanks for your contribution!

I think we should take a step back before digging into the current design.

A few questions based on the PR description:

Why checkpoint IcebergWriteResult metadata per write task? The checkpoint of .meta.pkl happens in the same task, on the same worker, at nearly the same time as the data file write. If the worker can fail, it can fail at any point — including during or before the metadata checkpoint write. So the metadata checkpoint doesn't provide a stronger durability guarantee than the data file write itself. This approach only covers the narrow case where the driver crashes after all workers fully complete but before on_write_complete() commits to the Iceberg catalog.
Can we just clean up uncommitted files instead?(Like rollback) If a job fails before the Iceberg catalog commit, the written data files are orphans — they exist on storage but aren't referenced by any snapshot, so they're invisible to queries. On retry, we can just use the existing checkpoint IDs (.parquet files) to skip already-committed rows, and either proactively delete orphaned files or let Iceberg's built-in remove_orphan_files handle cleanup.

This would be significantly simpler while still achieving the same correctness guarantees.

Curious to hear your thoughts — there may be context I'm missing that motivated this design.

WinkerDu · 2026-03-28T19:41:55Z

Hey @WinkerDu, thanks for your contribution!

I think we should take a step back before digging into the current design.

A few questions based on the PR description:

Why checkpoint IcebergWriteResult metadata per write task? The checkpoint of .meta.pkl happens in the same task, on the same worker, at nearly the same time as the data file write. If the worker can fail, it can fail at any point — including during or before the metadata checkpoint write. So the metadata checkpoint doesn't provide a stronger durability guarantee than the data file write itself. This approach only covers the narrow case where the driver crashes after all workers fully complete but before on_write_complete() commits to the Iceberg catalog.

Can we just clean up uncommitted files instead?(Like rollback) If a job fails before the Iceberg catalog commit, the written data files are orphans — they exist on storage but aren't referenced by any snapshot, so they're invisible to queries. On retry, we can just use the existing checkpoint IDs (.parquet files) to skip already-committed rows, and either proactively delete orphaned files or let Iceberg's built-in remove_orphan_files handle cleanup.

This would be significantly simpler while still achieving the same correctness guarantees.

Curious to hear your thoughts — there may be context I'm missing that motivated this design.

@owenowenisme Thank you for the reply.

Iceberg writes are a two-step protocol: workers write data files to storage, then the driver performs a catalog commit that references those files in a new snapshot. Until the catalog commit happens, files are orphans and are invisible to queries.
Checkpoint IDs only capture which input rows were processed, not which Iceberg data files were produced (i.e., the DataFile objects / file metadata required for the commit).
In the crash-before-commit scenario, it is possible to have:
- data files already written to storage, and
- checkpoint IDs already persisted,
- but no completed Iceberg catalog commit.
On retry, if we “just use existing checkpoint IDs to skip rows”, we will not reprocess those rows, which means we also will not regenerate the corresponding DataFile metadata for them.
Without the DataFile metadata, the driver has nothing to commit for the skipped rows, so those previously written orphan files remain unreferenced by any snapshot and never become visible. The job can succeed while the table is missing data, i.e., silent data loss.
Therefore, under the existing Iceberg write + commit semantics, “skip already-processed checkpoint IDs” by itself is not sufficient: skipping work must be paired with a way to recover the exact set of data files that should be committed.

WinkerDu · 2026-03-30T04:00:50Z

@owenowenisme Did you just leave a comment and edit or delete it afterward? I couldn’t find it in the PR. :-)

owenowenisme

I don't think we need to modify the datasource/datasink to support checkpointing here. When we added checkpoint support for parquet (PRs #59409 and #61821), the parquet datasource and datasink were left completely untouched — the only change was a minor filename API refactor to make filenames deterministic. All checkpoint logic lives in planner/checkpoint/ and checkpoint/, implemented as pre/post-write transforms that wrap around the existing write operator. The datasinks remain entirely unaware of checkpointing.

This PR breaks that pattern by adding checkpoint-specific code into three places:

IcebergDatasink.write() — builds per-block write results and stores them in ctx.kwargs
IcebergDatasink.on_write_complete() — loads .meta.pkl checkpoint files and merges with write returns
IcebergDatasource.get_read_tasks() — filters plan files using a checkpoint partition filter

All of this can be done without modifying the datasource/datasink:

Read-side filtering: plan_read_op.py already injects a post-read transform that filters out checkpointed row IDs — no need to touch get_read_tasks().
Write-side checkpointing: plan_write_op.py already injects pre/post-write transforms. The IcebergWriteResult (with DataFile paths) is available in write_returns after the write completes — the post-write transform can extract what it needs from there.
Recovery in on_write_complete: This can be handled by a checkpoint-aware wrapper at the planner level that merges recovered results into the WriteResult before calling the original on_write_complete. The datasink just sees a complete WriteResult and does its normal schema reconciliation and catalog commit.

Checkpoint logic should stay where it belongs — datasources and datasinks should be unaware of its existence.

owenowenisme · 2026-03-30T09:58:07Z

@@ -264,6 +264,86 @@ def _preprocess_data_pipeline(
        return checkpoint_ds.sort(self.id_column)


+class IcebergCheckpointLoader:


CAn we create a new file (something like iceberg_checkpoint.py) and move iceberg related logic there and maybe try to extend the base checkpoint mechanism? The checkpoint filter here should be general.

We can also move iceberg_datasource.py & iceberg_datasink.py under the folder _internal/datasource/iceberg

owenowenisme · 2026-03-30T10:06:52Z

@@ -171,21 +174,63 @@ def write_block_checkpoint(self, block: BlockAccessor):
        if block.num_rows() == 0:
            return

-        file_name = f"{uuid.uuid4()}.parquet"
+        file_id = str(uuid.uuid4())


Same problem here, we dont want to make the checkpoint_writer iceberg specific, try to extend on top of it

owenowenisme · 2026-03-30T15:07:18Z

+            # Serialize first so pickling errors don't occur after the parquet file
+            # has already been persisted.
+            #
+            # Write metadata before parquet to avoid a "parquet-only" partial state:


What if the error happen after .meta.pkl is created and before or during the checkpoint parquet is written? I think it will give us duplicate pkl file, why not just fuse this behavior with the current checkpoint 2pc?

richardliaw · 2026-04-30T00:40:13Z

@WinkerDu - any new update on this?

WinkerDu · 2026-04-30T03:41:04Z

@WinkerDu - any new update on this?

@richardliaw
I'll update the patch today

owenowenisme · 2026-05-06T13:29:39Z

@WinkerDu Any update on this? Thanks!

Resolved conflicts in: - python/ray/data/_internal/planner/planner.py: adopt pure branch's callbacks list pattern while removing stale supports_ckpt local variable - python/ray/data/checkpoint/checkpoint_filter.py: merge iceberg backend support from HEAD with 2-phase commit cleanup and numpy-based filter architecture from pure branch; unified CheckpointManager.__init__ to retain ckpt_config, checkpoint_path_unwrapped, and partition_filter fields Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Change-Id: Id3985c9242ea1feddc70b6e277985a4cf87c6db3

Change-Id: I9a5d27a66c6911307c1ab30f0679c4130ae3282f

Change-Id: Ie4f4c36d7f08f712aca963f887750e53292967ad

Change-Id: I736b092ef4ced0f63639be7abfe8f08e942d792f

Change-Id: If7e588fc469ed10e682df02958d0419ce1f31977

Change-Id: I0c362009433beb3c33e3e7d4d1e75319b078d61e

Change-Id: Ie7cbde9620fd63dc2d214077bafcac374d46933f

cursor

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

^{Reviewed by Cursor Bugbot for commit 2afc08a. Configure here.}

cursor · 2026-05-11T10:59:55Z

+        if checkpoint_config is not None:
+            checkpoint_config.delete_checkpoint_on_success = _original_delete_on_success
+        _delete_checkpoint_after_commit()
+        return result


Duplicate data on re-run with retained checkpoint metadata

Medium Severity

When delete_checkpoint_on_success is False, .meta.pkl files from a successful run persist. On the next pipeline execution, merge_recovered_iceberg_write_results loads these old results and re-commits their data files to the Iceberg table, because the new run's write_returns have different file paths so deduplication doesn't filter them out. This produces duplicate rows in the table for any non-first successful run when checkpoints are retained.

Additional Locations (1)

python/ray/data/checkpoint/iceberg_checkpoint.py#L139-L218

^{Reviewed by Cursor Bugbot for commit 2afc08a. Configure here.}

WinkerDu · 2026-05-13T07:44:16Z

Hi, @owenowenisme @richardliaw this pr has been updated.

This PR implements checkpointing for Iceberg table read/write operations based on Ray’s 2-phase commit workflow.
The read/write flow for Iceberg commit intermediate data (.meta.pkl) has been adapted to the existing 2-phase commit workflow.

coolderli · 2026-08-17T07:13:51Z

@WinkerDu @owenowenisme any new update about this? Thanks.

*I used CODEX to analyze this problem and create this PR. I've reviewed the code and tests and stand by them. This summary is written completely by a human (me) other than very light copy editing by an LLM.* The problem I ran into was that with 50000 Iceberg read tasks the driver was serializing 50000 copies of the same information as it serialized the tasks into its object store. This PR replaces those copies with an `ObjectRef` to the same information in Ray's distributed object store and the workers fetch that data when needed. This ended up being a meaningful performance gain in driver-side task construction and serialization at scale, see the benchmarking section below. In my production use case it reduced the Iceberg read-task construction and serialization from over an hour to less than one minute. CODEX generated summary follows: ----- Large Iceberg scans currently serialize the same PyIceberg scan state into every Ray `ReadTask`. On a production-shaped scan with 50,000 read tasks, that repeated metadata prevented Ray from finishing task construction, delayed autoscaling indefinitely, and projected far beyond the head node's object-store capacity. This change stores the immutable scan state once in Ray's object store. Each read task now carries only its file chunk and the shared `ObjectRef`; the worker resolves that reference before calling the existing Iceberg read implementation. Read semantics, planning, projection, filtering, and delete handling are unchanged. ## Why this is not duplicate work [#49054](#49054) fixed [#49107](#49107) by removing the datasource's `self` reference from each read closure. The replacement partial still captures table IO, full table metadata, the row filter, and the projected schema in every task. Searches of open PRs for that issue, `IcebergDatasource.get_read_tasks`, Iceberg `ObjectRef`, `table_metadata`, serialized Iceberg reads, and Iceberg object-store usage found no PR addressing that residual duplication. [#61753](#61753) touches this datasource for unrelated checkpoint filtering. ## Benchmark A controlled single-node benchmark isolated the driver-side work changed by this PR. It created a real local Iceberg table with 1,028 fields, projected 290 fields, and replicated a planned FileScanTask to produce different task counts. For each count, the benchmark constructed the ReadTasks and stored each with ray.put, matching plan_read_op.py. Object sizes were measured through get_local_object_locations. The patched totals include both the serialized ReadTasks and the single shared-state object. | Read tasks | Base object-store bytes | Patched object-store bytes | Reduction | | ---: | ---: | ---: | ---: | | 100 | 30.45 MB | 12.35 MB | 59.5% | | 1,000 | 304.54 MB | 122.40 MB | 59.8% | | 5,000 | 1.523 GB | 611.52 MB | 59.8% | At 1,000 tasks, median task construction and serialization time across three runs fell from 2.55 seconds to 0.46 seconds, a 5.5x improvement. The local fixture’s shared state was 117 KB, substantially smaller than the production workload’s metadata. It therefore demonstrates the scaling behavior without reproducing the production-sized state. On the production-shaped scan, the serialized size of one read task fell from approximately 1.63 MiB to 106 KiB. The patched implementation constructed all 50,000 tasks in 50.10 seconds, submitted the reads, and allowed autoscaling to begin. The stock implementation did not finish task construction or submit a read. Iceberg planning remained approximately 100 seconds in both cases, as expected; this change targets task serialization and construction rather than file planning. ## Tests ```text .venv/bin/pytest -q python/ray/data/tests/datasource/test_iceberg.py --tb=short # 55 passed PATH="$PWD/.venv/bin:/opt/homebrew/bin:/usr/bin:/bin" bazel test //python/ray/data:test_iceberg \ --test_output=errors --nocache_test_results --action_env=PATH \ --test_env=VIRTUAL_ENV="$PWD/.venv" # PASSED pre-commit run --files \ python/ray/data/_internal/datasource/iceberg_datasource.py \ python/ray/data/tests/datasource/test_iceberg.py # All applicable hooks passed ``` AI assistance was used to investigate the bottleneck, prototype the change, implement the patch, and draft tests and this description. This PR is a draft pending the human submitter's line-by-line review before review is requested. --------- Signed-off-by: Aaron Niskode-Dossett <aniskodedossett@etsy.com>

owenowenisme · 2026-09-11T02:27:22Z

@WinkerDu Could you resolve the conflicts?

WinkerDu added 7 commits February 26, 2026 02:25

[Data] Add Iceberg-backend in Checkpointing (ray-project#59870)

23ec4be

Change-Id: Iee31177fe923a779b958ad18aeea738adbf273f7

fix notimplementerror raise

34dd591

Change-Id: Ib43c5ef4160fba7b567e8b970edad09568f9aa48

fix lint

0e410f3

Change-Id: Ib148e684b6e193a307b0a3ddd02d5627ac17209a

fix cursor bot

1c5e51c

Change-Id: I058a20dad17a1d1653befc7d52ecf29709a53e58

fix TOCTOU race

4fc5f14

Change-Id: Iab693a43fbe0ff3e07f42a4039c6e76b232a1e32

fix cursor cr

d11fe61

Change-Id: Ie91b8235fd9d61b1521e204b66241d395b9a47ce

fix parallel commit

24dd4a0

Change-Id: Ib0fed415eb864e140c6aaee82d584d945f0c7402

WinkerDu requested a review from a team as a code owner March 15, 2026 16:36

gemini-code-assist Bot reviewed Mar 15, 2026

View reviewed changes

WinkerDu mentioned this pull request Mar 15, 2026

[Data] Add Iceberg-backend in Checkpointing #61316

Closed

cursor Bot reviewed Mar 15, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

Comment thread python/ray/data/checkpoint/checkpoint_filter.py Outdated

WinkerDu force-pushed the master-iceberg-checkpoint-3 branch from 10735a0 to e933ed1 Compare March 15, 2026 19:03

ray-gardener Bot added data Ray Data-related issues community-contribution Contributed by the community labels Mar 15, 2026

cursor Bot reviewed Mar 15, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

WinkerDu force-pushed the master-iceberg-checkpoint-3 branch from e933ed1 to 04f2433 Compare March 15, 2026 20:05

cursor Bot reviewed Mar 15, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

Comment thread python/ray/data/checkpoint/checkpoint_writer.py Outdated

WinkerDu force-pushed the master-iceberg-checkpoint-3 branch from 04f2433 to a1bf03f Compare March 19, 2026 09:21

cursor Bot reviewed Mar 19, 2026

View reviewed changes

Comment thread python/ray/data/_internal/planner/checkpoint/plan_write_op.py Outdated

cursor Bot reviewed Mar 20, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

cursor Bot reviewed Mar 20, 2026

View reviewed changes

Comment thread python/ray/data/checkpoint/checkpoint_filter.py Outdated

cursor Bot reviewed Mar 22, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

cursor Bot reviewed Mar 22, 2026

View reviewed changes

Comment thread python/ray/data/checkpoint/checkpoint_filter.py Outdated

cursor Bot reviewed Mar 24, 2026

View reviewed changes

Comment thread python/ray/data/_internal/datasource/iceberg_datasink.py Outdated

owenowenisme self-assigned this Mar 24, 2026

owenowenisme reviewed Mar 27, 2026

View reviewed changes

owenowenisme reviewed Mar 30, 2026

View reviewed changes

WinkerDu and others added 2 commits May 7, 2026 17:43

fix iceberg checkpoint merge issue & fix 2-phase commit issure

e121564

Change-Id: I9a5d27a66c6911307c1ab30f0679c4130ae3282f

WinkerDu force-pushed the master-iceberg-checkpoint-3 branch from 10adb61 to e121564 Compare May 8, 2026 15:34

cursor Bot reviewed May 8, 2026

View reviewed changes

Comment thread python/ray/data/checkpoint/checkpoint_filter.py Outdated

Comment thread python/ray/data/datasource/file_datasink.py Outdated

cursor Bot reviewed May 9, 2026

View reviewed changes

Comment thread python/ray/data/_internal/planner/checkpoint/plan_write_op.py

Comment thread python/ray/data/_internal/planner/planner.py

update ut

2eec93c

Change-Id: Ie4f4c36d7f08f712aca963f887750e53292967ad

WinkerDu force-pushed the master-iceberg-checkpoint-3 branch from 4ae9720 to 2eec93c Compare May 11, 2026 03:06

cursor Bot reviewed May 11, 2026

View reviewed changes

Comment thread python/ray/data/checkpoint/checkpoint_filter.py

Comment thread python/ray/data/checkpoint/checkpoint_filter.py Outdated

Comment thread python/ray/data/_internal/execution/execution_callback.py Outdated

fix ai cr

936c932

Change-Id: I736b092ef4ced0f63639be7abfe8f08e942d792f

cursor Bot reviewed May 11, 2026

View reviewed changes

Comment thread python/ray/data/tests/test_checkpoint_for_iceberg.py Outdated

remove unused ut

a9ff0c9

Change-Id: If7e588fc469ed10e682df02958d0419ce1f31977

cursor Bot reviewed May 11, 2026

View reviewed changes

Comment thread python/ray/data/checkpoint/util.py Outdated

WinkerDu and others added 3 commits May 11, 2026 17:14

remove unused util

377c687

Change-Id: I0c362009433beb3c33e3e7d4d1e75319b078d61e

Merge branch 'master' into master-iceberg-checkpoint-3

3e099c8

fix nit

2afc08a

Change-Id: Ie7cbde9620fd63dc2d214077bafcac374d46933f

cursor Bot reviewed May 11, 2026

View reviewed changes

my-vegetable-has-exploded mentioned this pull request Jun 28, 2026

[Data] Support Lance in Checkpointing #59871

Open

This was referenced Jul 16, 2026

[Data] Deduplicate Iceberg read task state dossett/ray#1

Draft

[Data] Deduplicate Iceberg read task state #64811

Merged

bveeramani added this to the Data issue and PR backlog milestone Aug 19, 2026

		@@ -264,6 +264,86 @@ def _preprocess_data_pipeline(
		return checkpoint_ds.sort(self.id_column)


		class IcebergCheckpointLoader:

Uh oh!

Conversation

WinkerDu commented Mar 15, 2026

Description

Related issues

Additional information

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

WinkerDu commented Mar 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

owenowenisme left a comment

Choose a reason for hiding this comment

Uh oh!

WinkerDu commented Mar 28, 2026

Uh oh!

WinkerDu commented Mar 30, 2026

Uh oh!

owenowenisme left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

owenowenisme Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

owenowenisme Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

owenowenisme Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

richardliaw commented Apr 30, 2026

Uh oh!

WinkerDu commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

owenowenisme commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor Bot May 11, 2026

Choose a reason for hiding this comment

Duplicate data on re-run with retained checkpoint metadata

Uh oh!

WinkerDu commented May 13, 2026

Uh oh!

coolderli commented Aug 17, 2026

Uh oh!

owenowenisme commented Sep 11, 2026

Uh oh!

Reviewers

WinkerDu commented Apr 30, 2026 •

edited

Loading

owenowenisme commented May 6, 2026 •

edited

Loading