ray-project--ray/python/ray/data/tests/test_checkpoint.py

import csv
import os
import random
from typing import List, Literal, Union

import numpy as np
import pandas as pd
import pyarrow
import pyarrow as pa
import pyarrow.parquet as pq
import pytest
from pyarrow.fs import FileSelector, FileType, LocalFileSystem
from pytest_lazy_fixtures import lf as lazy_fixture

import ray
from ray.data._internal.arrow_ops import transform_pyarrow
from ray.data._internal.datasource.csv_datasource import CSVDatasource
from ray.data._internal.datasource.parquet_datasink import ParquetDatasink
from ray.data._internal.execution.interfaces import TaskContext
from ray.data._internal.execution.operators.map_operator import MapOperator
from ray.data._internal.logical.interfaces.logical_plan import LogicalPlan
from ray.data._internal.logical.operators import Read, Write
from ray.data._internal.logical.optimizers import get_execution_plan
from ray.data._internal.planner.checkpoint.plan_read_op import (
    _get_checkpoint_map_transformer,
    plan_read_op_with_checkpoint_filter,
)
from ray.data._internal.planner.checkpoint.plan_write_op import (
    WRITE_UUID_KWARG_NAME,
    _generate_base_filename,
    _generate_prepare_checkpoint_transform,
)
from ray.data.block import BlockAccessor
from ray.data.checkpoint import CheckpointConfig
from ray.data.checkpoint.checkpoint_filter import (
    IdColumnCheckpointManager,
    NumpyArrayBasedCheckpointFilter,
)
from ray.data.checkpoint.checkpoint_writer import (
    PENDING_CHECKPOINT_SUFFIX,
    BatchBasedCheckpointWriter,
)
from ray.data.checkpoint.interfaces import (
    CheckpointBackend,
    InvalidCheckpointingConfig,
)
from ray.data.checkpoint.util import PrefixTrie
from ray.data.context import DataContext
from ray.data.datasource import BlockBasedFileDatasink, RowBasedFileDatasink
from ray.data.datasource.path_util import _unwrap_protocol
from ray.data.tests.conftest import *  # noqa
from ray.tests.conftest import *  # noqa

# User-provided ID column name
ID_COL = "id"
# Number of rows in the sample data
SAMPLE_DATA_NUM_ROWS = 10

# Auto-use `restore_data_context` for each test and apply 300-second timeout to all tests.
pytestmark = [
    pytest.mark.usefixtures("restore_data_context"),
    pytest.mark.timeout(300),
]


@pytest.fixture
def generate_sample_data_csv(tmp_path):
    def _generate(id_type: Literal["int", "str"] = "int"):
        # Generate a dummy dataset with SAMPLE_DATA_NUM_ROWS rows and columns [ID_COL, "col1"]
        ids = (
            [f"id_{i}" for i in range(SAMPLE_DATA_NUM_ROWS)]
            if id_type == "str"
            else list(range(SAMPLE_DATA_NUM_ROWS))
        )

        data = [{ID_COL: id_val, "col1": random.random()} for id_val in ids]

        f_path = os.path.join(tmp_path, "sample_data.csv")
        with open(f_path, mode="w", newline="") as file:
            writer = csv.DictWriter(file, fieldnames=data[0].keys())
            writer.writeheader()
            writer.writerows(data)
        return f_path

    return _generate


@pytest.fixture
def checkpoint_path(tmp_path):
    """Fixture to provide a temporary checkpoint path."""
    return str(tmp_path / "checkpoint")


@pytest.fixture
def data_output_path(data_path):
    """Fixture to provide a standardized data output path."""
    return os.path.join(data_path, "output")


@pytest.fixture
def generate_sample_data_parquet(tmp_path):
    def _generate():
        f_dir = os.path.join(tmp_path, "sample_data_parquet")
        os.makedirs(f_dir, exist_ok=True)
        # Generate a dummy dataset with SAMPLE_DATA_NUM_ROWS rows and columns [ID_COL, "col1"]
        df = pd.DataFrame(
            [{ID_COL: i, "col1": random.random()} for i in range(SAMPLE_DATA_NUM_ROWS)]
        )

        f_path = os.path.join(f_dir, "sample_data.parquet")
        # Write 3 row groups per file with uneven distribution of rows per row group
        table = pa.table(df)
        row_group_size = max(1, SAMPLE_DATA_NUM_ROWS // 3)
        pq.write_table(table, f_path, row_group_size=row_group_size)
        return f_dir

    return _generate


@pytest.fixture
def generate_sample_physical_plan(generate_sample_data_csv, tmp_path):
    ctx = ray.data.DataContext.get_current()

    datasource = CSVDatasource(generate_sample_data_csv())

    read_op = Read(datasource, datasource, -1, None)
    write_path = os.path.join(tmp_path, "output")
    write_op = Write(ParquetDatasink(write_path), input_dependencies=[read_op])
    logical_plan = LogicalPlan(write_op, ctx)
    physical_plan, _ = get_execution_plan(logical_plan)
    yield physical_plan


def _get_batch_based_files(ckpt_path: str, fs) -> List[str]:
    """Get checkpoint file paths for batch-based backends."""
    if fs is None:
        if not os.path.exists(ckpt_path):
            return []
        return [os.path.join(ckpt_path, f) for f in os.listdir(ckpt_path)]
    else:
        files = fs.get_file_info(
            FileSelector(_unwrap_protocol(ckpt_path), allow_not_found=True)
        )
        return [file_info.path for file_info in files if file_info.is_file]


def _read_batch_file_ids(file_paths: List[str], id_column: str, fs) -> List[int]:
    """Read IDs from batch-based checkpoint files."""
    ids = []
    for file_path in file_paths:
        if fs is None:
            table = pa.parquet.read_table(file_path)
        else:
            with fs.open_input_file(file_path) as f:
                table = pa.parquet.read_table(f)
        df = table.to_pandas()
        ids.extend(df[id_column].tolist())
    return ids


def read_ids_from_checkpoint_files(config: CheckpointConfig) -> List[Union[int, str]]:
    """Reads the checkpoint files and returns a sorted list of IDs which have been checkpointed."""
    # Batch-based backends
    if config.backend in (
        CheckpointBackend.FILE_STORAGE,
        CheckpointBackend.CLOUD_OBJECT_STORAGE,
    ):
        file_paths = _get_batch_based_files(config.checkpoint_path, config.filesystem)
        return sorted(
            _read_batch_file_ids(file_paths, config.id_column, config.filesystem)
        )

    else:
        raise ValueError(f"Invalid backend: {config.backend}")


class TestCheckpointConfig:
    @pytest.mark.parametrize("id_column", ["", 1])
    def test_invalid_id_column(self, id_column, local_path):
        with pytest.raises(
            InvalidCheckpointingConfig,
            match="Checkpoint ID column",
        ):
            CheckpointConfig(id_column, local_path)

    def test_override_backend_emits_deprecation_warning(self):
        with pytest.warns(FutureWarning, match="deprecated"):
            CheckpointConfig(
                ID_COL,
                "s3://bucket/path",
                override_backend=CheckpointBackend.FILE_STORAGE,
            )

    def test_default_checkpoint_path(self, s3_path, monkeypatch):
        with pytest.raises(
            InvalidCheckpointingConfig,
            match="CheckpointConfig.checkpoint_path",
        ):
            CheckpointConfig(ID_COL, None)

        default_bucket = s3_path
        monkeypatch.setenv(
            CheckpointConfig.DEFAULT_CHECKPOINT_PATH_BUCKET_ENV_VAR, default_bucket
        )

        config = CheckpointConfig(ID_COL, None)
        assert (
            config.checkpoint_path
            == f"{default_bucket}/{CheckpointConfig.DEFAULT_CHECKPOINT_PATH_DIR}"
        )

    @pytest.mark.parametrize("checkpoint_path", ["tmp/", "s3:/tmp", "s4://tmp"])
    def test_invalid_checkpoint_path(self, checkpoint_path):
        with pytest.raises(
            InvalidCheckpointingConfig,
            match="Invalid checkpoint path",
        ):
            CheckpointConfig(ID_COL, checkpoint_path)

    @pytest.mark.parametrize(
        "checkpoint_path",
        [
            lazy_fixture("local_path"),
            lazy_fixture("s3_path"),
        ],
    )
    def test_infer_filesystem_and_backend(self, checkpoint_path):
        config = CheckpointConfig(ID_COL, checkpoint_path)
        if checkpoint_path.startswith("/"):
            assert isinstance(config.filesystem, pyarrow.fs.LocalFileSystem)
            assert config.backend == CheckpointBackend.FILE_STORAGE
        else:
            assert isinstance(config.filesystem, pyarrow.fs.S3FileSystem)
            assert config.backend == CheckpointBackend.CLOUD_OBJECT_STORAGE

    @pytest.mark.parametrize(
        "checkpoint_path,fs,backend",
        [
            (
                lazy_fixture("local_path"),
                lazy_fixture("local_fs"),
                CheckpointBackend.FILE_STORAGE,
            ),
            (
                lazy_fixture("s3_path"),
                lazy_fixture("s3_fs"),
                CheckpointBackend.FILE_STORAGE,
            ),
            (
                lazy_fixture("local_path"),
                lazy_fixture("local_fs"),
                CheckpointBackend.CLOUD_OBJECT_STORAGE,
            ),
            (
                lazy_fixture("s3_path"),
                lazy_fixture("s3_fs"),
                CheckpointBackend.CLOUD_OBJECT_STORAGE,
            ),
        ],
    )
    def test_override_filesystem_and_backend(self, checkpoint_path, fs, backend):
        config = CheckpointConfig(
            ID_COL, checkpoint_path, override_filesystem=fs, override_backend=backend
        )
        assert config.filesystem is fs
        assert config.backend is backend

    def test_skip_inference_with_overrides(self):
        """Test that filesystem inference is skipped when override is provided."""
        # Inferring filesystem will fail if the path doesn't exist.
        path = "s3://non-existing-bucket/"
        fs = pyarrow.fs.S3FileSystem()
        config = CheckpointConfig(
            ID_COL,
            path,
            override_filesystem=fs,
        )
        assert config.filesystem is fs
        assert config.backend is CheckpointBackend.CLOUD_OBJECT_STORAGE


@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
        (
            CheckpointBackend.FILE_STORAGE,
            lazy_fixture("local_fs"),
            lazy_fixture("local_path"),
        ),
        (
            CheckpointBackend.CLOUD_OBJECT_STORAGE,
            lazy_fixture("s3_fs"),
            lazy_fixture("s3_path"),
        ),
    ],
)
def test_checkpoint_end_to_end(
    ray_start_10_cpus_shared,
    generate_sample_data_csv,
    backend,
    fs,
    data_path,
    data_output_path,
):
    """The end-to-end test for checkpoint."""

    class TestActor:
        def __init__(self):
            pass

        def __call__(self, batch):
            return batch

    ctx = ray.data.DataContext.get_current()
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        override_filesystem=fs,
        override_backend=backend,
    )

    csv_file = generate_sample_data_csv()

    # Generate checkpoint file
    checkpointed_ids = list(range(SAMPLE_DATA_NUM_ROWS // 2))
    expected_remaining_ids = sorted(
        set(range(SAMPLE_DATA_NUM_ROWS)) - set(checkpointed_ids)
    )

    ckpt_unwrapped = _unwrap_protocol(ckpt_path)
    if fs is None:
        os.makedirs(ckpt_unwrapped, exist_ok=True)
        ckpt_table = pa.table({ID_COL: checkpointed_ids})
        pq.write_table(
            ckpt_table, os.path.join(ckpt_unwrapped, "pre_checkpoint.parquet")
        )
    else:
        fs.create_dir(ckpt_unwrapped)
        ckpt_table = pa.table({ID_COL: checkpointed_ids})
        ckpt_file_path = os.path.join(ckpt_unwrapped, "pre_checkpoint.parquet")
        with fs.open_output_stream(ckpt_file_path) as f:
            pq.write_table(ckpt_table, f)

    ds = ray.data.read_csv(csv_file)

    # Execute the dataset with checkpointing enabled.
    ds = ds.map_batches(TestActor, concurrency=1)
    ds.write_parquet(data_output_path, filesystem=fs)

    # When execution succeeds, checkpoint data should be automatically deleted.
    # Check that the checkpoint directory is empty or doesn't exist
    if ctx.checkpoint_config.delete_checkpoint_on_success:
        try:
            unwrapped_path = _unwrap_protocol(ckpt_path)
            # Try to get file info for the checkpoint directory
            files = ctx.checkpoint_config.filesystem.get_file_info(
                pyarrow.fs.FileSelector(unwrapped_path, recursive=True)
            )
            # If we can get file info, the directory exists and should be empty
            assert (
                len(files) == 0
            ), f"Checkpoint directory should be empty but contains {len(files)} files"
        except (FileNotFoundError, OSError):
            # If directory doesn't exist, that's also fine (cleanup worked)
            pass

    # Ensure that the written data only contains the non-checkpointed rows.
    # Disable checkpointing before reading back to avoid filtering.
    ctx.checkpoint_config = None
    ds_readback = ray.data.read_parquet(data_output_path, filesystem=fs)
    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
    assert actual_output == expected_remaining_ids, (
        f"Expected only non-checkpointed IDs {expected_remaining_ids}, "
        f"but got {actual_output}"
    )


@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
        (
            CheckpointBackend.FILE_STORAGE,
            lazy_fixture("local_fs"),
            lazy_fixture("local_path"),
        ),
        (
            CheckpointBackend.CLOUD_OBJECT_STORAGE,
            lazy_fixture("s3_fs"),
            lazy_fixture("s3_path"),
        ),
    ],
)
def test_full_dataset_executed_for_non_write(
    ray_start_10_cpus_shared,
    generate_sample_data_parquet,
    backend,
    fs,
    data_path,
    data_output_path,
):
    """Tests that for an already fully checkpointed Dataset,
    calling `schema()` and `count()` should not skip checkpointing
    and should execute the full Dataset to get the correct information.
    """

    ctx = ray.data.DataContext.get_current()
    ctx.default_hash_shuffle_parallelism = 1
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        override_filesystem=fs,
        override_backend=backend,
    )

    parquet_dir = generate_sample_data_parquet()

    ds = ray.data.read_parquet(parquet_dir)

    ds = ds.map(lambda row: row)

    # Get the schema and count prior to writing the dataset.
    schema_before_write = ds.schema()
    count_before_write = ds.count()

    ds.write_parquet(data_output_path, filesystem=fs)

    # Recreate the same dataset, so that it will skip checkpointed rows.
    ds2 = ray.data.read_parquet(parquet_dir)
    ds2 = ds2.map(lambda row: row)

    # Check that when re-running a dataset which has already been completely
    # checkpointed, it does not skip any rows during `schema()` and `count()` calls.
    assert ds2.schema() == schema_before_write
    assert ds2.count() == count_before_write


@pytest.mark.parametrize(
    "ds_factory",
    [
        (lazy_fixture("generate_sample_data_parquet")),
        (lazy_fixture("generate_sample_data_parquet")),
    ],
)
@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
        (
            CheckpointBackend.FILE_STORAGE,
            lazy_fixture("local_fs"),
            lazy_fixture("local_path"),
        ),
        (
            CheckpointBackend.CLOUD_OBJECT_STORAGE,
            lazy_fixture("s3_fs"),
            lazy_fixture("s3_path"),
        ),
    ],
)
def test_recovery_no_missing_rows(
    ray_start_10_cpus_shared,
    ds_factory,
    backend,
    fs,
    data_path,
    data_output_path,
):
    """Tests that recovery is at_least_once: no missing rows after retry."""

    ctx = ray.data.DataContext.get_current()
    ctx.execution_options.preserve_order = True
    ctx.default_hash_shuffle_parallelism = 1
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")

    # Ensure checkpoint directory exists
    os.makedirs(ckpt_path, exist_ok=True)

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        override_filesystem=fs,
        override_backend=backend,
    )

    # Catch the custom TestException raised by FailActor.
    ctx.raise_original_map_exception = True

    @ray.remote(num_cpus=0)
    class Coordinator:
        def __init__(self):
            self._should_fail = True

        def disable_failure(self):
            self._should_fail = False

        def should_fail(self):
            return self._should_fail

    coordinator_actor = Coordinator.remote()

    class TestException(Exception):
        pass

    class FailActor:
        """Simple passthrough actor, which fails after a certain number of rows."""

        def __init__(self, coordinator_actor, max_num_items, checkpoint_config):
            self._should_fail = ray.get(coordinator_actor.should_fail.remote())
            self._max_num_items = max_num_items
            self._checkpoint_config = checkpoint_config

        def __call__(self, batch):
            # Get the ID column name from the checkpoint config
            id_col = self._checkpoint_config.id_column

            # Process each row in the batch
            ids = batch[id_col]

            for _, id in enumerate(ids):
                if self._should_fail and id == 2:
                    raise TestException(f"FailActor: Failing on row {id}")

            return batch

    # Use the ds_factory to create the dataset
    local_data_path = ds_factory()
    ds = ray.data.read_parquet(local_data_path)

    # Get the actual number of items from the dataset
    max_num_items = ds.count()

    ds = ds.map_batches(
        FailActor,
        fn_constructor_args=[coordinator_actor, max_num_items, ctx.checkpoint_config],
        concurrency=1,
        batch_size=None,
        num_cpus=1.1,  # Use a different num_cpus to avoid operator fusion.
    )

    # Should fail in the middle.
    with pytest.raises(TestException):
        ds.write_parquet(data_output_path, filesystem=fs, concurrency=1)

    ray.get(coordinator_actor.disable_failure.remote())
    # When executing the same dataset again, this should skip the already
    # checkpointed rows.
    ds.write_parquet(data_output_path, filesystem=fs, concurrency=1)

    # When execution succeeds, checkpoint data should be automatically deleted.
    assert read_ids_from_checkpoint_files(ctx.checkpoint_config) == []

    # Get the ID column name from the checkpoint config
    id_col = ctx.checkpoint_config.id_column

    # Disable checkpointing prior to reading back the data, so we don't skip any rows.
    ctx.checkpoint_config = None

    # Ensure that the written data is correct.
    ds_readback = ray.data.read_parquet(data_output_path, filesystem=fs)

    # Checkpointing is at_least_once, so duplicates are allowed; ensure no missing IDs.
    actual_output = [row[id_col] for row in ds_readback.iter_rows()]
    expected_output = list(range(max_num_items))
    assert set(actual_output) == set(expected_output)


@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
        (
            CheckpointBackend.FILE_STORAGE,
            lazy_fixture("local_fs"),
            lazy_fixture("local_path"),
        ),
        (
            CheckpointBackend.CLOUD_OBJECT_STORAGE,
            lazy_fixture("s3_fs"),
            lazy_fixture("s3_path"),
        ),
    ],
)
def test_manual_checkpoint_filters_ids(
    ray_start_10_cpus_shared,
    generate_sample_data_parquet,
    backend,
    fs,
    data_path,
    data_output_path,
):
    """Manually write checkpoint IDs and verify they are filtered from output."""
    ctx = ray.data.DataContext.get_current()
    ctx.default_hash_shuffle_parallelism = 1
    ckpt_path = os.path.join(data_path, "test_manual_checkpoint_files")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        override_filesystem=fs,
        override_backend=backend,
    )

    checkpoint_ids = [1, 3, 7]
    df = pd.DataFrame({ID_COL: checkpoint_ids})
    checkpoint_writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    checkpoint_writer.write_block_checkpoint(BlockAccessor.for_block(df))
    assert set(read_ids_from_checkpoint_files(ctx.checkpoint_config)) == set(
        checkpoint_ids
    )

    parquet_dir = generate_sample_data_parquet()
    ds = ray.data.read_parquet(parquet_dir)
    ds.write_parquet(data_output_path, filesystem=fs)

    # Disable checkpointing prior to reading back the data, so we don't skip any rows.
    ctx.checkpoint_config = None
    ds_readback = ray.data.read_parquet(data_output_path, filesystem=fs)
    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
    expected_output = sorted(
        [i for i in range(SAMPLE_DATA_NUM_ROWS) if i not in set(checkpoint_ids)]
    )
    assert actual_output == expected_output


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_pending_checkpoint_write_and_commit(fs, base_path):
    """Test the two-phase commit (2PC) checkpoint write and commit workflow.

    This verifies that:
    1. write_pending_checkpoint() creates a checkpoint file with PENDING suffix
    2. The pending checkpoint exists but the committed version does not
    3. commit_checkpoint() renames the pending file to committed (removes suffix)
    4. After commit, the pending file no longer exists and committed file exists
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    fs.create_dir(_unwrap_protocol(checkpoint_path))
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    df = pd.DataFrame({ID_COL: [1, 2], "col1": [0.1, 0.2]})

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id="test_000000_000000",
    )
    assert pending is not None
    assert fs.get_file_info(pending.pending_path).type != FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type == FileType.NotFound

    writer.commit_checkpoint(pending)
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type != FileType.NotFound

    # Read back and verify contents
    with fs.open_input_file(pending.committed_path) as f:
        table = pq.read_table(f)
    assert table.column(ID_COL).to_pylist() == [1, 2]


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_commit_checkpoint_idempotent_already_committed(fs, base_path):
    """Test that commit_checkpoint is idempotent when already committed.

    If the committed file already exists and pending doesn't, calling
    commit_checkpoint should succeed without error (no-op).
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    fs.create_dir(_unwrap_protocol(checkpoint_path))
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    df = pd.DataFrame({ID_COL: [1, 2]})

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id="test_000000_000000",
    )
    assert pending is not None

    # First commit succeeds
    writer.commit_checkpoint(pending)
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type != FileType.NotFound

    # Second commit should be idempotent (no error)
    writer.commit_checkpoint(pending)
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type != FileType.NotFound


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_commit_checkpoint_idempotent_both_exist(fs, base_path):
    """Test that commit_checkpoint cleans up when both files exist.

    If both committed and pending files exist (edge case), the pending
    file should be deleted and committed file preserved.
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    fs.create_dir(_unwrap_protocol(checkpoint_path))
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    df = pd.DataFrame({ID_COL: [1, 2]})

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id="test_000000_000000",
    )
    assert pending is not None

    # Commit normally
    writer.commit_checkpoint(pending)

    # Manually recreate the pending file to simulate edge case
    with fs.open_output_stream(pending.pending_path) as f:
        f.write(b"dummy")

    assert fs.get_file_info(pending.pending_path).type != FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type != FileType.NotFound

    # Commit should clean up the pending file
    writer.commit_checkpoint(pending)
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type != FileType.NotFound


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_commit_checkpoint_neither_exists(fs, base_path):
    """Test that commit_checkpoint raises error when neither file exists."""
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    fs.create_dir(_unwrap_protocol(checkpoint_path))
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    df = pd.DataFrame({ID_COL: [1, 2]})

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id="test_000000_000000",
    )
    assert pending is not None

    # Delete the pending file to simulate missing state
    fs.delete_file(pending.pending_path)
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound
    assert fs.get_file_info(pending.committed_path).type == FileType.NotFound

    # Commit should raise FileNotFoundError
    with pytest.raises(FileNotFoundError):
        writer.commit_checkpoint(pending)


@pytest.mark.parametrize("data_file_exists", [True, False])
@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_clean_pending_checkpoint(
    ray_start_10_cpus_shared, fs, base_path, data_file_exists
):
    """Test pending checkpoint cleanup removes incomplete writes.

    When a write fails after creating a pending checkpoint but before commit,
    the cleanup process must:
    1. Find pending checkpoint files, build prefix trie from their basenames
    2. Delete associated data files matching the prefix (if they exist)
    3. Delete the pending checkpoint files

    This test verifies cleanup works correctly whether the data file was
    actually written (data_file_exists=True) or not (data_file_exists=False).
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    data_dir = os.path.join(base_path, "data")
    for p in [checkpoint_path, data_dir]:
        fs.create_dir(_unwrap_protocol(p))

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)

    # Write a pending checkpoint (simulating pre-write phase)
    df = pd.DataFrame({ID_COL: [0]})
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id="test_000000_000000",
    )
    assert pending is not None
    assert fs.get_file_info(pending.pending_path).type != FileType.NotFound

    # Optionally create a data file matching the pending checkpoint prefix
    data_file = os.path.join(_unwrap_protocol(data_dir), "test_000000_000000.csv")
    if data_file_exists:
        with fs.open_output_stream(data_file) as f:
            f.write(b"id\n0\n")

    checkpoint_manager = IdColumnCheckpointManager(ctx.checkpoint_config, ctx)
    checkpoint_manager._clean_pending_checkpoints(data_dir, fs)

    # Data file should be deleted (if it existed)
    if data_file_exists:
        assert fs.get_file_info(data_file).type == FileType.NotFound
    # Pending checkpoint should be deleted
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_clean_pending_checkpoint_with_partitioned_data(
    ray_start_10_cpus_shared, fs, base_path
):
    """Test pending checkpoint cleanup removes files in partition subdirectories.

    When using ParquetDatasink with partition_cols, data files are written to
    subdirectories like output/col=val/file.parquet. The cleanup process must
    recursively search subdirectories to find and delete data files matching
    pending checkpoint prefixes.
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    data_dir = os.path.join(base_path, "data")
    for p in [checkpoint_path, data_dir]:
        fs.create_dir(_unwrap_protocol(p))

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)

    # Write a pending checkpoint for "pending_task" (simulating failed write)
    base_filename = "pending_task"
    df = pd.DataFrame({ID_COL: [0, 1, 2]})
    id_column_data = BlockAccessor.for_block(df).to_arrow()[ID_COL]
    pending = writer.write_pending_checkpoint(
        id_column_data,
        checkpoint_id=base_filename,
    )
    assert pending is not None

    # Create data files matching the pending checkpoint in partition subdirectories
    partition_dirs = ["partition_col=a", "partition_col=b"]
    pending_data_files = []
    for partition_dir in partition_dirs:
        partition_path = os.path.join(_unwrap_protocol(data_dir), partition_dir)
        fs.create_dir(partition_path)
        data_file = os.path.join(partition_path, f"{base_filename}-0.parquet")
        with fs.open_output_stream(data_file) as f:
            f.write(b"dummy data")
        pending_data_files.append(data_file)

    # Also create a matching file directly in base dir
    base_data_file = os.path.join(
        _unwrap_protocol(data_dir), f"{base_filename}-0.parquet"
    )
    with fs.open_output_stream(base_data_file) as f:
        f.write(b"dummy data")
    pending_data_files.append(base_data_file)

    # Create unrelated data files that should NOT be deleted
    unrelated_files = []
    for partition_dir in partition_dirs:
        partition_path = os.path.join(_unwrap_protocol(data_dir), partition_dir)
        data_file = os.path.join(partition_path, "other_task-0.parquet")
        with fs.open_output_stream(data_file) as f:
            f.write(b"dummy data")
        unrelated_files.append(data_file)

    # Verify all files were created
    for f in pending_data_files + unrelated_files:
        assert (
            fs.get_file_info(f).type != FileType.NotFound
        ), f"Expected file to exist: {f}"

    # Run cleanup
    checkpoint_manager = IdColumnCheckpointManager(ctx.checkpoint_config, ctx)
    checkpoint_manager._clean_pending_checkpoints(data_dir, fs)

    # Verify data files matching pending checkpoint were deleted
    for f in pending_data_files:
        assert (
            fs.get_file_info(f).type == FileType.NotFound
        ), f"Expected pending data file to be deleted: {f}"

    # Verify unrelated data files still exist
    for f in unrelated_files:
        assert (
            fs.get_file_info(f).type != FileType.NotFound
        ), f"Expected unrelated file to still exist: {f}"

    # Verify pending checkpoint was also deleted
    assert fs.get_file_info(pending.pending_path).type == FileType.NotFound


def test_clean_pending_checkpoints_nonexistent_path(ray_start_10_cpus_shared, tmp_path):
    """Test that _clean_pending_checkpoints handles a non-existent checkpoint dir.

    On cloud storage like Azure Blob Storage, listing a non-existent directory
    raises FileNotFoundError (unlike S3 which returns an empty result). The
    FileSelector uses allow_not_found=True so that first-run scenarios where
    the checkpoint directory does not yet exist succeed without error.
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(tmp_path, "nonexistent", "deeply", "nested", "path")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
    )

    checkpoint_manager = IdColumnCheckpointManager(ctx.checkpoint_config, ctx)

    data_dir = os.path.join(tmp_path, "data")
    os.makedirs(data_dir, exist_ok=True)

    # Should succeed gracefully (no pending checkpoints to clean)
    checkpoint_manager._clean_pending_checkpoints(data_dir)


def test_prepare_checkpoint_transform_writes_pending(tmp_path):
    """Test that the pre-write checkpoint transform writes a pending checkpoint.

    This verifies that:
    1. _generate_prepare_checkpoint_transform() creates a transform that writes
       a pending checkpoint file
    2. The pending checkpoint filename matches the base filename with .pending suffix
    """

    class MockFileDatasink(BlockBasedFileDatasink):
        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
            file.write(b"")

    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(tmp_path, "checkpoint")
    os.makedirs(checkpoint_path, exist_ok=True)
    data_output_path = os.path.join(tmp_path, "output")
    os.makedirs(data_output_path, exist_ok=True)
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
    )

    datasink = MockFileDatasink(data_output_path, file_format="csv")
    checkpoint_writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    transform = _generate_prepare_checkpoint_transform(ctx, datasink, checkpoint_writer)

    ctx_task = TaskContext(task_idx=0, op_name="test")
    ctx_task.kwargs[WRITE_UUID_KWARG_NAME] = "test-write-uuid"

    df = pd.DataFrame({ID_COL: [0], "col1": [0.1]})
    list(transform._apply_transform(ctx_task, [df]))

    pending_files = [
        f
        for f in os.listdir(checkpoint_path)
        if f.endswith(f"{PENDING_CHECKPOINT_SUFFIX}.parquet")
    ]
    assert len(pending_files) == 1

    # Verify pending checkpoint filename matches the base filename
    base_filename = _generate_base_filename(datasink, ctx_task)
    assert pending_files[0] == f"{base_filename}{PENDING_CHECKPOINT_SUFFIX}.parquet"


def test_2pc_fail_retry_cleans_pending_checkpoints(
    ray_start_10_cpus_shared,
    tmp_path,
):
    """Test end-to-end 2PC cleanup: fail during write, retry succeeds after cleanup.

    This is an integration test for the full two-phase commit cleanup flow:
    1. First write attempt: pre-write creates pending checkpoint, datasink
       writes partial data then fails, leaving pending checkpoint files
    2. Retry: the checkpoint filter's _clean_pending_checkpoints() detects
       pending files, deletes them along with partial data files
    3. Second write attempt: succeeds, writing complete data with no duplicates

    Verifies that after cleanup, all pending checkpoints are removed and
    the final output contains the correct data.
    """
    ctx = ray.data.DataContext.get_current()
    ctx.raise_original_map_exception = True
    ctx.default_hash_shuffle_parallelism = 1

    checkpoint_path = os.path.join(tmp_path, "checkpoint")
    data_output_path = os.path.join(tmp_path, "output")
    os.makedirs(checkpoint_path, exist_ok=True)

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
    )

    @ray.remote(num_cpus=0)
    class FailController:
        def __init__(self):
            self._should_fail = True

        def should_fail(self):
            return self._should_fail

        def disable_failure(self):
            self._should_fail = False

    controller = FailController.remote()

    class FailOnceCSVDatasink(BlockBasedFileDatasink):
        def __init__(self, path: str, controller):
            super().__init__(path, file_format="csv")
            self._controller = controller

        def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
            if ray.get(self._controller.should_fail.remote()):
                # Write a partial file and then fail to simulate incomplete write.
                file.write(b"id\n0\n")
                raise RuntimeError("Simulated write failure")
            block.to_pandas().to_csv(file, index=False)

    datasink = FailOnceCSVDatasink(data_output_path, controller)
    ds = ray.data.range(SAMPLE_DATA_NUM_ROWS, override_num_blocks=1)

    with pytest.raises(RuntimeError, match="Simulated write failure"):
        ds.write_datasink(datasink, ray_remote_args={"max_retries": 0})

    # After failure, there should be pending checkpoint files (written in
    # pre-write phase, before the data write failed)
    pending_files = [
        f
        for f in os.listdir(checkpoint_path)
        if f.endswith(f"{PENDING_CHECKPOINT_SUFFIX}.parquet")
    ]
    assert pending_files, "Expected pending checkpoint files after failed write."

    ray.get(controller.disable_failure.remote())
    ds.write_datasink(datasink, ray_remote_args={"max_retries": 0})

    # After successful retry, pending checkpoints should be cleaned up
    pending_files_after = [
        f
        for f in os.listdir(checkpoint_path)
        if f.endswith(f"{PENDING_CHECKPOINT_SUFFIX}.parquet")
    ]
    assert pending_files_after == []

    ctx.checkpoint_config = None
    ds_readback = ray.data.read_csv(data_output_path)
    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
    expected_output = sorted(range(SAMPLE_DATA_NUM_ROWS))
    assert actual_output == expected_output


@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
        (
            CheckpointBackend.FILE_STORAGE,
            lazy_fixture("local_fs"),
            lazy_fixture("local_path"),
        ),
        (
            CheckpointBackend.CLOUD_OBJECT_STORAGE,
            lazy_fixture("s3_fs"),
            lazy_fixture("s3_path"),
        ),
    ],
)
def test_skip_checkpoint_flag(
    ray_start_10_cpus_shared,
    generate_sample_data_csv,
    backend,
    fs,
    data_path,
):
    """Test that for a valid Dataset with checkpointing enabled, calling methods like
    `schema()` and `count()` should skip checkpointing and not create any checkpoint
    files. Subsequently calling `write_xxx()` on the same dataset should have
    checkpointing enabled."""

    ctx = ray.data.DataContext.get_current()
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")
    ctx.checkpoint_config = CheckpointConfig(
        ID_COL,
        ckpt_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
        override_backend=backend,
    )

    def generate_ds():
        ds = ray.data.read_csv(generate_sample_data_csv())

        ds = ds.map(lambda row: row)
        return ds

    ds = generate_ds()

    # Calling `ds.schema()` should skip checkpointing.
    assert ds.schema() is not None
    assert len(read_ids_from_checkpoint_files(ctx.checkpoint_config)) == 0

    # Calling `ds.count()` should skip checkpointing.
    ds = generate_ds()
    assert ds.count() is not None
    assert len(read_ids_from_checkpoint_files(ctx.checkpoint_config)) == 0

    # Calling `ds.write_xxx()` afterwards should enable checkpointing.
    ds.write_parquet(os.path.join(data_path, "output"), filesystem=fs)

    # Check what checkpoint files exist
    checkpoint_files = read_ids_from_checkpoint_files(ctx.checkpoint_config)

    assert len(checkpoint_files) == SAMPLE_DATA_NUM_ROWS


def test_checkpoint_with_missing_id_column(
    ray_start_10_cpus_shared,
    generate_sample_data_csv,
    tmp_path,
):
    """Test that checkpointing fails gracefully when the configured id_column doesn't exist in the data."""

    ctx = ray.data.DataContext.get_current()
    ckpt_path = os.path.join(tmp_path, "test_checkpoint_output_files")
    # Configure checkpointing with an id_column that doesn't exist in the CSV data
    ctx.checkpoint_config = CheckpointConfig(
        id_column="nonexistent_column",
        checkpoint_path=ckpt_path,
        delete_checkpoint_on_success=False,
    )

    def generate_ds():
        ds = ray.data.read_csv(generate_sample_data_csv())
        ds = ds.map(lambda row: row)
        return ds

    ds = generate_ds()

    # The write operation should fail because the id_column doesn't exist
    with pytest.raises(
        ValueError,
        match="ID column nonexistent_column is absent in the block to be written",
    ):
        ds.write_parquet(os.path.join(tmp_path, "output"))


def test_dict_checkpoint_config(checkpoint_path):
    """Test that a dict checkpoint config can be used to create a CheckpointConfig."""
    context = ray.data.DataContext.get_current()
    fs = LocalFileSystem()
    context.checkpoint_config = {
        "id_column": ID_COL,
        "checkpoint_path": checkpoint_path,
        "override_filesystem": fs,
        "override_backend": "CLOUD_OBJECT_STORAGE",
    }
    assert context.checkpoint_config.id_column == ID_COL
    assert context.checkpoint_config.checkpoint_path == checkpoint_path
    assert context.checkpoint_config.filesystem is fs
    assert context.checkpoint_config.backend == CheckpointBackend.CLOUD_OBJECT_STORAGE


def test_write_block_checkpoint_with_pandas_df(restore_data_context, tmp_path):
    ctx = ray.data.DataContext.get_current()

    ctx.checkpoint_config = CheckpointConfig(
        ID_COL,
        str(tmp_path),
    )
    df = pd.DataFrame({ID_COL: [0, 1]})
    expected_ids = [0, 1]

    checkpoint_writer = BatchBasedCheckpointWriter(ctx.checkpoint_config)
    checkpoint_writer.write_block_checkpoint(BlockAccessor.for_block(df))

    assert len(os.listdir(tmp_path)) == 1
    checkpoint_filename = os.listdir(tmp_path)[0]
    checkpoint_path = tmp_path / checkpoint_filename

    table = pa.parquet.read_table(checkpoint_path)
    df = table.to_pandas()
    written_ids = df[ID_COL].tolist()
    assert written_ids == expected_ids


def test_filter_rows_for_block():
    """Test NumpyArrayBasedCheckpointFilter.filter_rows_for_block."""

    # Common test setup
    checkpoint_path = "/mock/path"

    # Test with simple ID column
    config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
    )

    # Create a mock block.
    block = pyarrow.table(
        {
            ID_COL: list(range(10)),
            "data": [str(i) for i in range(10)],
        }
    )
    # Create a mock checkpointed_ids with multiple chunks.
    chunk1 = pyarrow.table({ID_COL: [1, 2, 4]})
    chunk2 = pyarrow.table({ID_COL: [6, 8, 9, 11]})
    chunk3 = pyarrow.table({ID_COL: [12, 13]})
    checkpointed_ids = pyarrow.concat_tables([chunk1, chunk2, chunk3])
    assert len(checkpointed_ids[ID_COL].chunks) == 3

    checkpoint_ids_array = []
    for ckpt_chunk in checkpointed_ids[ID_COL].chunks:
        checkpoint_ids_array.append(
            transform_pyarrow.to_numpy(ckpt_chunk, zero_copy_only=False)
        )
    checkpointed_ids_ndarray = np.concatenate(checkpoint_ids_array)
    checkpointed_ids_ref = ray.put(checkpointed_ids_ndarray)

    expected_block = pyarrow.table(
        {
            ID_COL: [0, 3, 5, 7],
            "data": ["0", "3", "5", "7"],
        }
    )

    filter_instance = NumpyArrayBasedCheckpointFilter(config, checkpointed_ids_ref)
    filtered_block = filter_instance.filter_rows_for_block(
        block=block,
    )

    assert filtered_block.equals(expected_block)


def test_checkpoint_restore_after_full_execution(
    ray_start_10_cpus_shared,
    tmp_path,
    generate_sample_data_parquet,
    checkpoint_path,
):
    """Test checkpoint restore after full execution of data pipeline. This is
    done by retaining the checkpoint metadata files with
    delete_checkpoint_on_success=False.
    """

    def run_simple_pipeline(
        checkpoint_config: CheckpointConfig, input_path: str, output_path: str
    ) -> int:
        """Run a simple pipeline with checkpointing."""
        from ray.data.datasource import WriteResult

        ctx = DataContext.get_current()
        ctx.checkpoint_config = checkpoint_config
        ctx.default_hash_shuffle_parallelism = 1
        ds = ray.data.read_parquet(input_path)

        # Patch `on_write_complete` to get the WriteResult.
        num_rows_written = None
        original_on_write_complete = ParquetDatasink.on_write_complete

        def patched_on_write_complete(self, write_result: WriteResult[None]):
            nonlocal num_rows_written
            num_rows_written = write_result.num_rows
            return original_on_write_complete(self, write_result)

        ParquetDatasink.on_write_complete = patched_on_write_complete

        ds.write_parquet(output_path)
        return int(num_rows_written)

    # Create test paths
    input_data_path = generate_sample_data_parquet()
    data_output_path = str(tmp_path / "output")

    # Create checkpoint config
    checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        override_backend=CheckpointBackend.FILE_STORAGE,
        delete_checkpoint_on_success=False,
    )

    # First run: create checkpoint
    num_rows_first = run_simple_pipeline(
        checkpoint_config, input_data_path, data_output_path
    )
    assert (
        num_rows_first == SAMPLE_DATA_NUM_ROWS
    ), f"Expected {SAMPLE_DATA_NUM_ROWS} rows, got {num_rows_first}"

    # Check if checkpoint files were created
    assert os.path.exists(checkpoint_path), "No checkpoint directory created!"

    # Second run: should use checkpoint
    num_rows_second = run_simple_pipeline(
        checkpoint_config, input_data_path, data_output_path
    )
    assert (
        num_rows_second == 0  # No rows should be written
    ), f"Expected 0 rows, got {num_rows_second}"


@pytest.mark.parametrize(
    "data_path",
    [
        (lazy_fixture("local_path")),
    ],
)
def test_checkpoint_map_transformer(
    ray_start_10_cpus_shared,
    data_path,
):
    """Test checkpoint map transformer."""
    ctx = ray.data.DataContext.get_current()
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL, checkpoint_path=ckpt_path
    )

    checkpointed_ids_ndarray = np.array([1, 3, 5, 7, 9], dtype=np.int64)
    checkpointed_ids_ref = ray.put(checkpointed_ids_ndarray)

    map_transformer = _get_checkpoint_map_transformer(ctx, checkpointed_ids_ref)
    map_transformer.init()

    block = pyarrow.table(
        {
            ID_COL: list(range(10)),
            "data": [str(i) for i in range(10)],
        }
    )
    filtered_blocks = map_transformer.apply_transform(
        input_blocks=[block],
        ctx=TaskContext(task_idx=0, op_name="test_checkpoint"),
    )

    filtered_block = next(iter(filtered_blocks))
    expected_block = pyarrow.table(
        {
            ID_COL: [0, 2, 4, 6, 8],
            "data": ["0", "2", "4", "6", "8"],
        }
    )
    assert filtered_block.equals(expected_block)


@pytest.mark.parametrize(
    "data_path",
    [
        (lazy_fixture("local_path")),
    ],
)
def test_plan_read_op_with_checkpoint_filter_no_checkpoint_dir(
    ray_start_10_cpus_shared, generate_sample_data_csv, data_path
):
    """Test that when checkpoint directory does not exist,
    plan_read_op_with_checkpoint_filter returns the original read physical operator."""
    ctx = ray.data.DataContext.get_current()

    csv_file = generate_sample_data_csv()
    datasource = CSVDatasource(csv_file)

    # checkpoint_path points to a non-existent directory
    ckpt_path = os.path.join(str(data_path), "non_existent_ckpt_dir")
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        delete_checkpoint_on_success=False,
    )

    read_op = Read(datasource, datasource, -1, None)
    physical_op = plan_read_op_with_checkpoint_filter(
        None,
        None,
        op=read_op,
        physical_children=[],
        data_context=ctx,
    )

    # Should return the original ReadCSV op
    assert physical_op.name == "ReadCSV"


@pytest.mark.parametrize(
    "data_path",
    [
        (lazy_fixture("local_path")),
    ],
)
def test_plan_read_op_with_checkpoint_filter_empty_checkpoint_dir(
    ray_start_10_cpus_shared, generate_sample_data_csv, data_path
):
    """Test that when checkpoint directory exists but is empty,
    plan_read_op_with_checkpoint_filter returns the original read physical operator."""
    ctx = ray.data.DataContext.get_current()

    csv_file = generate_sample_data_csv()
    datasource = CSVDatasource(csv_file)

    # Create an empty checkpoint directory
    ckpt_path = os.path.join(str(data_path), "empty_ckpt_dir")
    os.makedirs(ckpt_path, exist_ok=True)

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        delete_checkpoint_on_success=False,
    )

    read_op = Read(datasource, datasource, -1, None)
    physical_op = plan_read_op_with_checkpoint_filter(
        None,
        None,
        op=read_op,
        physical_children=[],
        data_context=ctx,
    )

    # Should return the original ReadCSV op
    assert physical_op.name == "ReadCSV"


@pytest.mark.parametrize(
    "data_path",
    [
        (lazy_fixture("local_path")),
    ],
)
def test_plan_read_op_with_checkpoint_filter_with_valid_checkpoint(
    ray_start_10_cpus_shared,
    generate_sample_data_csv,
    data_path,
):
    """Test that when a valid checkpoint exists,
    plan_read_op_with_checkpoint_filter returns a CheckpointFilter MapOperator."""
    ctx = ray.data.DataContext.get_current()

    csv_file = generate_sample_data_csv()
    datasource = CSVDatasource(csv_file)

    # Create a checkpoint directory with valid checkpoint data
    ckpt_path = os.path.join(str(data_path), "valid_ckpt_dir")
    os.makedirs(ckpt_path, exist_ok=True)

    # Write some checkpoint IDs (e.g., IDs 0-4 are already processed)
    checkpointed_ids = pa.table({ID_COL: list(range(5))})
    pq.write_table(checkpointed_ids, os.path.join(ckpt_path, "ckpt_0.parquet"))

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        delete_checkpoint_on_success=False,
    )

    read_op = Read(datasource, datasource, -1, None)
    physical_op = plan_read_op_with_checkpoint_filter(
        None,
        None,
        op=read_op,
        physical_children=[],
        data_context=ctx,
    )

    # Should return a CheckpointFilter MapOperator
    assert isinstance(physical_op, MapOperator)
    assert physical_op.name == "CheckpointFilter"


@pytest.mark.parametrize(
    "backend,fs,data_path",
    [
        (CheckpointBackend.FILE_STORAGE, None, lazy_fixture("local_path")),
    ],
)
def test_checkpoint_with_string_typed_id(
    ray_start_10_cpus_shared,
    generate_sample_data_csv,
    backend,
    fs,
    data_path,
    data_output_path,
):
    """Test the checkpoint when the ID column is of type string."""

    class TestActor:
        def __init__(self):
            pass

        def __call__(self, batch):
            return batch

    ctx = ray.data.DataContext.get_current()
    ckpt_path = os.path.join(data_path, "test_checkpoint_output_files")

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=ckpt_path,
        override_filesystem=fs,
        override_backend=backend,
    )

    csv_file = generate_sample_data_csv(id_type="str")

    # Generate checkpoint file
    checkpointed_ids = [f"id_{id}" for id in range(SAMPLE_DATA_NUM_ROWS // 2)]
    expected_remaining_ids = sorted(
        {f"id_{id}" for id in range(SAMPLE_DATA_NUM_ROWS)} - set(checkpointed_ids)
    )

    ckpt_unwrapped = _unwrap_protocol(ckpt_path)

    os.makedirs(ckpt_unwrapped, exist_ok=True)
    ckpt_table = pa.table({ID_COL: checkpointed_ids})
    pq.write_table(ckpt_table, os.path.join(ckpt_unwrapped, "pre_checkpoint.parquet"))

    ds = ray.data.read_csv(csv_file)

    # Execute the dataset with checkpointing enabled.
    ds = ds.map_batches(TestActor, concurrency=1)
    ds.write_parquet(data_output_path, filesystem=fs)

    # Ensure that the written data only contains the non-checkpointed rows.
    # Disable checkpointing before reading back to avoid filtering.
    ctx.checkpoint_config = None
    ds_readback = ray.data.read_parquet(data_output_path, filesystem=fs)
    actual_output = sorted([row[ID_COL] for row in ds_readback.iter_rows()])
    assert actual_output == expected_remaining_ids, (
        f"Expected only non-checkpointed IDs {expected_remaining_ids}, "
        f"but got {actual_output}"
    )


class FailAfterWriteParquetDatasink(ParquetDatasink):
    """Test helper: ParquetDatasink that fails AFTER writing data (simulates post-write crash).

    This simulates the failure scenario where:
    - Data is successfully written to the output file
    - But the process crashes/fails before the checkpoint can be committed
    This is the critical case that 2PC is designed to handle - the data file
    exists but is "uncommitted" and should be cleaned up on recovery.
    """

    def __init__(self, path: str, fail_threshold: int = 100, **kwargs):
        super().__init__(path, **kwargs)
        self._fail_threshold = fail_threshold

    def write(self, blocks, ctx):
        blocks_list = list(blocks)

        # Check if any block has id > threshold
        should_fail = False
        for block in blocks_list:
            accessor = BlockAccessor.for_block(block)
            df = accessor.to_pandas()
            if ID_COL in df.columns and df[ID_COL].max() > self._fail_threshold:
                should_fail = True
                break

        # First, write the blocks normally
        result = super().write(iter(blocks_list), ctx)

        # Then fail if threshold exceeded (simulates post-write failure)
        if should_fail:
            raise RuntimeError(
                f"Simulated failure: block contains {ID_COL} > {self._fail_threshold}"
            )

        return result


def test_partial_failure_no_duplicates(
    ray_start_10_cpus_shared,
    tmp_path,
):
    """Test checkpoint deduplication: partial failure + retry produces no duplicate rows.

    This is the key correctness test for the checkpoint deduplication feature.
    It verifies that when a write pipeline fails partway through:
    1. Already-committed rows (from successful blocks before failure) are tracked
    2. Uncommitted rows (from blocks that failed after writing data) are cleaned up
    3. On retry, only uncommitted rows are re-written
    4. Final output has exactly the expected rows with NO duplicates

    The test uses run_tag to verify which rows came from which run, confirming
    that committed rows from run 1 are preserved and not re-written in run 2.

    Note: This test requires ray.shutdown() + ray.init() mid-test to flush
    in-flight checkpoint writes, which is incompatible with mock S3 (pyarrow's
    S3 subsystem gets finalized during shutdown). The checkpoint storage layer
    is already tested with S3 in other parameterized tests.
    """
    num_rows = 1000
    fail_threshold = 100

    # Create paths
    input_path = tmp_path / "input"
    output_path = tmp_path / "output"
    checkpoint_path_dir = tmp_path / "checkpoint"
    for path in [input_path, output_path, checkpoint_path_dir]:
        path.mkdir(exist_ok=True)

    # Create sample data (1000 rows with unique IDs)
    df = pd.DataFrame(
        {ID_COL: range(num_rows), "value": [f"row_{i}" for i in range(num_rows)]}
    )
    df.to_parquet(input_path / "data.parquet", index=False)

    # Configure checkpointing
    ctx = DataContext.get_current()
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=str(checkpoint_path_dir),
        delete_checkpoint_on_success=False,
    )

    def add_run_tag(batch, run_tag):
        """Add a run_tag column to identify which run wrote the data."""
        batch["run_tag"] = [run_tag] * len(batch[ID_COL])
        return batch

    # Run 1: Use the failing datasink - should write some blocks then fail
    with pytest.raises(RuntimeError, match="Simulated failure"):
        ds = ray.data.read_parquet(str(input_path))
        ds = ds.repartition(200)  # 200 blocks
        ds = ds.map_batches(lambda b: add_run_tag(b, "first"), batch_size=None)
        failing_datasink = FailAfterWriteParquetDatasink(
            str(output_path), fail_threshold=fail_threshold
        )
        ds.write_datasink(failing_datasink, ray_remote_args={"max_retries": 0})

    # Shutdown Ray to ensure all in-flight tasks complete before checking state.
    # This addresses the race condition where background tasks may still be writing
    # checkpoint files after the exception is raised.
    ray.shutdown()
    ray.init()

    # Run 2: Use regular write_parquet - should resume from checkpoint
    ds2 = ray.data.read_parquet(str(input_path))
    ds2 = ds2.repartition(200)
    ds2 = ds2.map_batches(lambda b: add_run_tag(b, "second"), batch_size=None)
    ds2.write_parquet(str(output_path))

    result = ray.data.read_parquet(str(output_path)).to_pandas()

    assert len(result) == num_rows, f"Expected {num_rows} rows, got {len(result)}"

    # Check for duplicates
    assert result[ID_COL].is_unique, (
        f"Duplicate IDs found: "
        f"{sorted(result[result.duplicated(ID_COL, keep=False)][ID_COL].unique().tolist())}"
    )

    # Verify that some rows came from first run (before failure) and rest from second
    run_tag_counts = result["run_tag"].value_counts()
    assert "first" in run_tag_counts.index, "Expected some rows from first run"
    assert "second" in run_tag_counts.index, "Expected some rows from second run"


def test_partial_failure_no_duplicates_partitioned(
    ray_start_10_cpus_shared,
    tmp_path,
):
    """Test checkpoint deduplication with multi-level partitioned parquet output.

    Same as test_partial_failure_no_duplicates, but writes partitioned output
    using 3 partition columns, creating deeply nested subdirectories
    (e.g., output/region=us/category=x/tier=1/file.parquet). This exercises
    the recovery path where data files must be found via recursive search
    using the data_file_dir passed through the call chain.

    Note: This test requires ray.shutdown() + ray.init() mid-test to flush
    in-flight checkpoint writes, which is incompatible with mock S3 (pyarrow's
    S3 subsystem gets finalized during shutdown). The checkpoint storage layer
    is already tested with S3 in other parameterized tests.
    """
    num_rows = 1000
    fail_threshold = 100

    # Create paths
    input_path = tmp_path / "input"
    output_path = tmp_path / "output"
    checkpoint_path_dir = tmp_path / "checkpoint"
    for path in [input_path, output_path, checkpoint_path_dir]:
        path.mkdir(exist_ok=True)

    # Create sample data with 3 partition columns for deeply nested output
    # (e.g., output/region=us/category=x/tier=1/file.parquet)
    regions = ["us", "eu"]
    categories = ["x", "y", "z"]
    tiers = [1, 2]
    df = pd.DataFrame(
        {
            ID_COL: range(num_rows),
            "value": [f"row_{i}" for i in range(num_rows)],
            "region": [regions[i % len(regions)] for i in range(num_rows)],
            "category": [categories[i % len(categories)] for i in range(num_rows)],
            "tier": [tiers[i % len(tiers)] for i in range(num_rows)],
        }
    )
    df.to_parquet(input_path / "data.parquet", index=False)

    # Configure checkpointing
    ctx = DataContext.get_current()
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=str(checkpoint_path_dir),
        delete_checkpoint_on_success=False,
    )

    def add_run_tag(batch, run_tag):
        """Add a run_tag column to identify which run wrote the data."""
        batch["run_tag"] = [run_tag] * len(batch[ID_COL])
        return batch

    partition_cols = ["region", "category", "tier"]

    # Run 1: Use the failing datasink with partition_cols
    with pytest.raises(RuntimeError, match="Simulated failure"):
        ds = ray.data.read_parquet(str(input_path))
        ds = ds.repartition(200)
        ds = ds.map_batches(lambda b: add_run_tag(b, "first"), batch_size=None)
        failing_datasink = FailAfterWriteParquetDatasink(
            str(output_path),
            fail_threshold=fail_threshold,
            partition_cols=partition_cols,
        )
        ds.write_datasink(failing_datasink, ray_remote_args={"max_retries": 0})

    # Shutdown Ray to ensure all in-flight tasks complete before checking state.
    ray.shutdown()
    ray.init()

    # Run 2: Use regular write_parquet with partition_cols - should resume
    ds2 = ray.data.read_parquet(str(input_path))
    ds2 = ds2.repartition(200)
    ds2 = ds2.map_batches(lambda b: add_run_tag(b, "second"), batch_size=None)
    ds2.write_parquet(str(output_path), partition_cols=partition_cols)

    result = ray.data.read_parquet(str(output_path)).to_pandas()

    assert len(result) == num_rows, f"Expected {num_rows} rows, got {len(result)}"

    # Check for duplicates
    assert result[ID_COL].is_unique, (
        f"Duplicate IDs found: "
        f"{sorted(result[result.duplicated(ID_COL, keep=False)][ID_COL].unique().tolist())}"
    )

    # Verify that some rows came from first run (before failure) and rest from second
    run_tag_counts = result["run_tag"].value_counts()
    assert "first" in run_tag_counts.index, "Expected some rows from first run"
    assert "second" in run_tag_counts.index, "Expected some rows from second run"

    # Verify partitioned output structure: partition subdirectories should exist
    output_subdirs = [
        d for d in os.listdir(str(output_path)) if os.path.isdir(output_path / d)
    ]
    assert len(output_subdirs) > 0, "Expected partition subdirectories in output"


class TextRowDatasink(RowBasedFileDatasink):
    """Test helper: RowBasedFileDatasink that writes each row as a text file."""

    def __init__(self, path: str, **kwargs):
        super().__init__(path, file_format="txt", **kwargs)

    def write_row_to_file(self, row, file):
        file.write(f"{row[ID_COL]},{row['value']}".encode())


class FailAfterWriteTextRowDatasink(TextRowDatasink):
    """Test helper: TextRowDatasink that fails AFTER writing data."""

    def __init__(self, path: str, fail_threshold: int = 100, **kwargs):
        super().__init__(path, **kwargs)
        self._fail_threshold = fail_threshold

    def write(self, blocks, ctx):
        blocks_list = list(blocks)

        should_fail = False
        for block in blocks_list:
            accessor = BlockAccessor.for_block(block)
            df = accessor.to_pandas()
            if ID_COL in df.columns and df[ID_COL].max() > self._fail_threshold:
                should_fail = True
                break

        result = super().write(iter(blocks_list), ctx)

        if should_fail:
            raise RuntimeError(
                f"Simulated failure: block contains {ID_COL} > {self._fail_threshold}"
            )

        return result


def test_partial_failure_no_duplicates_row_based(
    ray_start_10_cpus_shared,
    tmp_path,
):
    """Test checkpoint deduplication with row-based datasink (one file per row).

    Row-based datasinks write multiple files per task, e.g.:
      write_uuid_000000_000000_000000.txt
      write_uuid_000000_000000_000001.txt
    Recovery must match all of them via the checkpoint prefix (write_uuid_000000).
    This validates the PrefixTrie-based matching works for row-based writes.
    """
    num_rows = 1000
    fail_threshold = 100

    # Create paths
    input_path = tmp_path / "input"
    output_path = tmp_path / "output"
    checkpoint_path_dir = tmp_path / "checkpoint"
    for path in [input_path, output_path, checkpoint_path_dir]:
        path.mkdir(exist_ok=True)

    # Create sample data
    df = pd.DataFrame(
        {ID_COL: range(num_rows), "value": [f"row_{i}" for i in range(num_rows)]}
    )
    df.to_parquet(input_path / "data.parquet", index=False)

    # Configure checkpointing
    ctx = DataContext.get_current()
    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=str(checkpoint_path_dir),
        delete_checkpoint_on_success=False,
    )

    def add_run_tag(batch, run_tag):
        batch["run_tag"] = [run_tag] * len(batch[ID_COL])
        return batch

    # Run 1: Use the failing row-based datasink
    with pytest.raises(RuntimeError, match="Simulated failure"):
        ds = ray.data.read_parquet(str(input_path))
        ds = ds.repartition(200)
        ds = ds.map_batches(lambda b: add_run_tag(b, "first"), batch_size=None)
        failing_datasink = FailAfterWriteTextRowDatasink(
            str(output_path), fail_threshold=fail_threshold
        )
        ds.write_datasink(failing_datasink, ray_remote_args={"max_retries": 0})

    ray.shutdown()
    ray.init()

    # Run 2: Use the same row-based datasink (non-failing) to resume
    ds2 = ray.data.read_parquet(str(input_path))
    ds2 = ds2.repartition(200)
    ds2 = ds2.map_batches(lambda b: add_run_tag(b, "second"), batch_size=None)
    ds2.write_datasink(TextRowDatasink(str(output_path)))

    # Read all output text files and parse them
    output_files = [f for f in os.listdir(str(output_path)) if f.endswith(".txt")]
    rows = []
    for fname in output_files:
        with open(os.path.join(str(output_path), fname)) as f:
            content = f.read()
            id_val, value = content.split(",", 1)
            rows.append({ID_COL: int(id_val), "value": value})

    result = pd.DataFrame(rows)

    assert len(result) == num_rows, f"Expected {num_rows} rows, got {len(result)}"

    # Check for duplicates
    assert result[ID_COL].is_unique, (
        f"Duplicate IDs found: "
        f"{sorted(result[result.duplicated(ID_COL, keep=False)][ID_COL].unique().tolist())}"
    )


def test_prefix_trie():
    """Test PrefixTrie insert and has_prefix_of operations."""
    trie = PrefixTrie()
    trie.insert("abc")
    trie.insert("def")
    trie.insert("ab")

    # "ab" is a prefix of "abc", "abcd", "ab"
    assert trie.has_prefix_of("abc")
    assert trie.has_prefix_of("abcd")
    assert trie.has_prefix_of("ab")
    assert trie.has_prefix_of("abxyz")
    assert trie.has_prefix_of("def")
    assert trie.has_prefix_of("defgh")

    # No prefix matches for these
    assert not trie.has_prefix_of("a")
    assert not trie.has_prefix_of("xyz")
    assert not trie.has_prefix_of("d")
    assert not trie.has_prefix_of("de")
    assert not trie.has_prefix_of("")


def test_prefix_trie_empty():
    """Test that an empty PrefixTrie returns False for all queries."""
    trie = PrefixTrie()
    assert not trie.has_prefix_of("")
    assert not trie.has_prefix_of("abc")
    assert not trie.has_prefix_of("anything")


@pytest.mark.parametrize(
    "fs,base_path",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
    ],
    ids=["local", "s3"],
)
def test_clean_pending_checkpoints_no_pending(ray_start_10_cpus_shared, fs, base_path):
    """Test that cleanup is a no-op when there are no pending checkpoint files.

    With no pending checkpoints, no data files should be deleted. This is the
    normal case after a successful write where all checkpoints were committed.
    """
    ctx = ray.data.DataContext.get_current()
    checkpoint_path = os.path.join(base_path, "checkpoint")
    fs.create_dir(_unwrap_protocol(checkpoint_path))

    ctx.checkpoint_config = CheckpointConfig(
        id_column=ID_COL,
        checkpoint_path=checkpoint_path,
        delete_checkpoint_on_success=False,
        override_filesystem=fs,
    )

    # Create data files but no pending checkpoints
    data_dir = os.path.join(base_path, "data")
    fs.create_dir(_unwrap_protocol(data_dir))
    data_file = os.path.join(_unwrap_protocol(data_dir), "some_data.parquet")
    with fs.open_output_stream(data_file) as f:
        f.write(b"dummy")

    checkpoint_manager = IdColumnCheckpointManager(ctx.checkpoint_config, ctx)
    checkpoint_manager._clean_pending_checkpoints(data_dir, fs)

    # Data file should still exist (no pending checkpoints means nothing to clean)
    assert fs.get_file_info(data_file).type != FileType.NotFound


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))