def test_retry_with_missing_records(test_session):
    """Test retry functionality with missing records."""
    # Create source dataset
    source_ids = [1, 2, 3]
    source_contents = ["first", "second", "third"]

    dc.read_values(id=source_ids, content=source_contents, session=test_session).save(
        "source_data"
    )

    def simple_process(id: int, content: str, attempt: int) -> ProcessingResult:
        return ProcessingResult(
            processed_content=content.upper(),
            processed_at=datetime.now(tz=timezone.utc).isoformat(),
            error="",
            attempt=attempt,
        )

    # Process only first 2 records
    # Create partial result dataset (missing id=3)
    partial_result = (
        dc.read_dataset("source_data", session=test_session)
        .setup(attempt=lambda: 1)
        .filter(C("id") < 3)
        .map(result=simple_process)
        .save("partial_result")
    )

    assert partial_result.count() == 2

    # Use retry with delta_retry=True to process missing records
    retry_chain = (
        dc.read_dataset(
            "source_data",
            session=test_session,
            delta=True,
            delta_on="id",
            delta_retry=True,
        )
        .setup(attempt=lambda: 2)
        .map(result=simple_process)
        .save("partial_result")
    )

    # Should now have all 3 records
    assert retry_chain.count() == 3

    # Verify all records are present
    ids = set(retry_chain.collect("id"))
    assert ids == {1, 2, 3}

    final_first_attempts_count = retry_chain.filter(C("result.attempt") == 1).count()
    final_missing_attempts_count = retry_chain.filter(C("result.attempt") == 2).count()

    # Only missing records should have attempt 2
    assert final_missing_attempts_count == 1
    assert final_first_attempts_count == 2
Test retry functionality with missing records.
tests/func/test_retry.py