feat(checkpoint): support checkpoint on lance, stage Lance DataSink write results

[everySympathy]

Changes Made

This PR adds the Daft core plumbing needed for checkpoint-aware Python DataSink writers. Lance is the first consumer.

The design keeps Daft core format-agnostic. Daft core does not need to understand Lance fragments. Instead, a Python DataSink can opt in through a small hook and ask Daft core to stage its per-input write_results in the checkpoint store. The sink package can later read those staged write_results and interpret them in its own format-specific final commit logic.

Python DataSink
        |
        v
checkpoint_file_format()
        |
        v
Rust pipeline
        |
        v
BlockingSinkNode
        |
        v
stage write_results
        |
        v
checkpoint(id)
        |
        v
Python finalize() reads staged metadata

This PR adds:

• CheckpointFileFormat.Lance.
• FileFormat::Lance in checkpoint metadata format tags.
• A default Python DataSink.checkpoint_file_format() implementation that returns None.
• Rust pipeline support for checkpoint-aware Python DataSinks.
• BlockingSinkNode::with_checkpoint wiring for opt-in DataSinks.
• DataFrame.write_lance(..., checkpoint=...) API plumbing.
• Append-only validation for checkpointed Lance writes.
• A driver-side fast path for Lance retries:
  • if the sink finds the idempotence key in Lance transaction history, Daft skips source scanning and worker writes;
  • finalize([]) is still called so the sink can mark pending checkpoint ids committed.

The Rust/Python boundary is intentionally narrow:

• Rust stages opaque write_results in the checkpoint store.
• Rust tags the metadata as Lance.
• Python daft_lance owns Lance-specific recovery and final commit logic.

Why

Checkpointed sink recovery needs source keys and sink output metadata to share the same checkpoint boundary.

For Lance, workers write fragments before the final table commit. The final table commit happens later in driver-side finalize(). If a task fails and the job retries, Daft source checkpointing may skip source keys that were already processed. In that case, the retry run may not reproduce the write_results for those previously processed inputs.

Daft core therefore needs a generic way to persist Python DataSink write_results before sealing the checkpoint id.

Testing

Tests added in this PR cover:

• Python exposure of CheckpointFileFormat.Lance.
• S3 checkpoint store encode/decode support for Lance metadata.
• Default Python DataSink.checkpoint_file_format() behavior.
• write_lance(..., checkpoint=...) API plumbing.
• Append-only validation for checkpointed Lance writes.

This PR is paired with daft-engine/daft-lance#16, which adds Lance-specific tests for:

• staged write_results recovery;
• idempotence-key based retry skipping;
• recovery when Lance commit succeeds but checkpoint mark_committed() fails;
• recovery when checkpoint staging succeeds but Lance commit fails.

Manual validation:

• Ran checkpointed read_parquet -> write_lance on a Ray cluster.
• Used TOS/S3-compatible storage.
• Retried with the same checkpoint store and idempotence key.
• Verified committed checkpoints were reused.
• Verified Lance fragments were not appended twice.
• Verified final Lance dataset row count and checkpoint statuses.

Related Issues

#6967

Related PR

daft-engine/daft-lance#16

[greptile-apps]

Greptile Summary

This PR adds the Daft core plumbing needed for checkpoint-aware Python DataSink writers, with LanceDataSink as the first consumer. It keeps Daft core format-agnostic: a DataSink opts in via a new checkpoint_file_format() hook, Rust stages the opaque write_results blob in the checkpoint store tagged with the format, and the sink's Python finalize() handles format-specific recovery.

• Adds FileFormat::Lance / CheckpointFileFormat.Lance through the full Rust↔Python stack (types.rs, config.rs, __init__.pyi), with correct serde(rename_all = "lowercase") serialisation and round-trip tests in s3.rs.
• Wires checkpoint staging into BlockingSinkNode for DataSink plans and introduces a driver-side idempotence fast path in write_lance(checkpoint=...) that skips source scanning when Lance already has a transaction for the idempotence key.
• Enforces append-only validation for checkpointed Lance writes and adds unit tests for the new Python surface (CheckpointFileFormat.Lance exposure, default checkpoint_file_format() return value).

Confidence Score: 4/5

The Rust checkpoint plumbing and Lance variant additions are clean and well-tested; the two findings are non-blocking style and efficiency observations.

The checkpoint data-path changes in pipeline.rs are straightforward and correctly guarded by the #[cfg(feature = "python")] gate. The datasink_checkpoint_file_format function being called twice is a minor redundancy rather than a functional defect. The write_lance(checkpoint=...) fast path correctly short-circuits to finalize([]), and the append-only guard is in place. The only behavioural concern is that third-party DataSink authors have no documented set of valid strings to return from checkpoint_file_format(), which could surface as a confusing Rust-level error in the future.

The datasink_checkpoint_file_format helper in src/daft-local-execution/src/pipeline.rs is worth a second look for the double-invocation pattern. daft/io/sink.py should document which format strings are accepted by checkpoint_file_format().

Important Files Changed

Filename	Overview
daft/daft/init.pyi	Adds Lance variant to CheckpointFileFormat stub enum; straightforward pyi update.
daft/dataframe/dataframe.py	Adds checkpoint parameter to write_lance, enforces append-only guard, and introduces a driver-side fast path that calls sink.checkpoint_commit_exists() on the external daft_lance.LanceDataSink; correctness depends on daft_lance >= checkpoint-support version.
daft/io/sink.py	Adds default checkpoint_file_format() to DataSink; valid return strings are undocumented in the method contract.
src/daft-checkpoint/src/config.rs	Adds Lance variant to PyCheckpointFileFormat and its From<FileFormat> impl; exhaustive match ensures future variants are caught at compile time.
src/daft-checkpoint/src/impls/s3.rs	Adds Lance to encode/decode round-trip tests; serde(rename_all = "lowercase") correctly serialises the new variant as "lance", consistent with existing format tags.
src/daft-checkpoint/src/types.rs	Adds Lance variant to FileFormat; #[non_exhaustive] + exhaustive match in config.rs ensures compile-time safety for future additions.
src/daft-local-execution/src/pipeline.rs	Adds datasink_checkpoint_file_format helper and wires it into both translate_physical_plan_to_pipeline and physical_plan_to_pipeline; the Python method is called twice per pipeline build for checkpoint-enabled DataSinks.
tests/checkpoint/test_checkpoint_config.py	Adds test that CheckpointFileFormat.Lance is exposed; test is minimal but covers the Python exposure requirement.
tests/io/test_sink.py	Adds test verifying that the default checkpoint_file_format() returns None; straightforward and correct.

Sequence Diagram

sequenceDiagram
    participant User as Python caller
    participant DF as DataFrame.write_lance
    participant Sink as LanceDataSink
    participant Pipeline as Rust Pipeline
    participant CkptStore as CheckpointStore

    User->>DF: "write_lance(uri, mode=append, checkpoint=commit)"
    DF->>Sink: "LanceDataSink(uri, ..., checkpoint=commit)"
    DF->>Sink: checkpoint_commit_exists()?
    alt idempotence key already committed in Lance
        Sink-->>DF: True
        DF->>Sink: finalize([])
        Note over Sink: mark_committed() only
        Sink-->>DF: MicroPartition
        DF-->>User: DataFrame (fast path)
    else normal path
        Sink-->>DF: False
        DF->>Pipeline: write_sink(sink)
        Pipeline->>Sink: checkpoint_file_format()
        Sink-->>Pipeline: Lance
        Pipeline->>Pipeline: BlockingSinkNode.with_checkpoint(store, id_map, Lance)
        loop per input batch
            Pipeline->>Sink: write(micropartitions)
            Sink-->>Pipeline: write_results
            Pipeline->>CkptStore: stage_files(write_results, Lance)
        end
        Pipeline->>CkptStore: checkpoint(id)
        Pipeline->>Sink: finalize(write_results)
        Note over Sink: Lance table commit + mark_committed()
        Sink-->>Pipeline: MicroPartition
        Pipeline-->>User: DataFrame
    end

Loading

_{Reviews (1): Last reviewed commit: "Merge branch 'main' into daft-lance-chec..." | Re-trigger Greptile}

[greptile-apps]

Double Python call to checkpoint_file_format()

datasink_checkpoint_file_format is invoked twice for any checkpoint-enabled DataSink plan: once here inside translate_physical_plan_to_pipeline to compute is_write_sink, and a second time inside physical_plan_to_pipeline's DataSink arm to call node.with_checkpoint(...). Both calls acquire the GIL and round-trip into Python. The result is pure: it will always return the same value for the same sink object. The is_write_sink result could be threaded down into physical_plan_to_pipeline (or cached on DataSinkInfo) so the Python method is only called once per pipeline build.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Valid format strings not documented

The docstring says the return value is "a format" but does not enumerate what strings Daft accepts. Right now the only accepted value is "Lance" (or "lance") — any other non-None return causes a ValueError deep inside the Rust pipeline builder with a message like unsupported DataSink checkpoint file format: "Iceberg". Third-party sink authors have no way to know valid values from reading this method's contract. Document the accepted strings explicitly (e.g., Returns: "Lance" or None).

[everySympathy]

@rohitkulshreshtha Hi Rohit, could you help to review this PR? Thanks!