chore: import upstream snapshot with attribution

python/ray/data/checkpoint/__init__.py

@@ -0,0 +1,3 @@
+from .interfaces import CheckpointBackend, CheckpointConfig
+
+__all__ = ["CheckpointConfig", "CheckpointBackend"]

python/ray/data/checkpoint/checkpoint_filter.py

@@ -0,0 +1,444 @@
+import abc
+import logging
+import os
+import posixpath
+import sys
+import time
+from abc import abstractmethod
+from typing import List, Optional, Tuple
+
+import numpy as np
+import pyarrow
+from pyarrow.fs import FileSelector, FileType
+
+import ray
+from ray._common.retry import call_with_retry
+from ray.data._internal.arrow_ops import transform_pyarrow
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data.block import Block, BlockMetadata, Schema
+from ray.data.checkpoint import CheckpointConfig
+from ray.data.checkpoint.checkpoint_writer import PENDING_CHECKPOINT_SUFFIX
+from ray.data.checkpoint.util import build_pending_checkpoint_trie
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import _unwrap_protocol
+from ray.types import ObjectRef
+
+logger = logging.getLogger(__name__)
+
+# Retry configuration for checkpoint recovery operations.
+# These can be overridden via environment variables for testing or tuning.
+CHECKPOINT_RECOVERY_MAX_ATTEMPTS = int(
+    os.environ.get("RAY_DATA_CHECKPOINT_RECOVERY_MAX_ATTEMPTS", "3")
+)
+CHECKPOINT_RECOVERY_MAX_BACKOFF_S = int(
+    os.environ.get("RAY_DATA_CHECKPOINT_RECOVERY_MAX_BACKOFF_S", "8")
+)
+
+
+def _numpy_size(array: np.ndarray) -> int:
+    """Calculate the size of a numpy ndarray."""
+    total_size = array.nbytes
+    if array.dtype == object:
+        sample_count = 10**4
+
+        if len(array) <= sample_count:
+            for item in array.flat:
+                total_size += sys.getsizeof(item)
+        else:
+            sample_total_size = 0
+            for item in array[:sample_count].flat:
+                sample_total_size += sys.getsizeof(item)
+            total_size += int(sample_total_size / sample_count * len(array))
+    return total_size
+
+
+@ray.remote(num_cpus=0)
+def _clean_pending_checkpoints_task(
+    checkpoint_path_unwrapped: str,
+    checkpoint_filesystem: pyarrow.fs.FileSystem,
+    data_file_dir_unwrapped: str,
+    data_file_filesystem: pyarrow.fs.FileSystem,
+) -> int:
+    """Delete data files that have matching pending checkpoint files, then
+    delete the pending checkpoints.
+
+    This runs as a remote task to avoid blocking the driver during potentially
+    slow filesystem operations (especially on cloud storage like S3).
+
+    Algorithm:
+    1. List all files in checkpoint dir, find those ending with .pending.parquet
+    2. Build a PrefixTrie from their basenames (strip .pending.parquet)
+    3. List all data files in data_file_dir (recursively for partitions)
+    4. For each data file, if trie.has_prefix_of(basename) -> delete it
+    5. Delete all the pending checkpoint files
+    6. Return count of pending checkpoints cleaned
+
+    Args:
+        checkpoint_path_unwrapped: The unwrapped checkpoint path.
+        checkpoint_filesystem: The filesystem for checkpoint files.
+        data_file_dir_unwrapped: The unwrapped directory where data files are
+            written (protocol prefix already stripped).
+        data_file_filesystem: The filesystem for data files. May differ from
+            checkpoint_filesystem (e.g., checkpoints on local disk, data on S3).
+
+    Returns:
+        Number of pending checkpoints cleaned.
+    """
+
+    def _clean() -> int:
+        # 1. List all files in checkpoint dir, find pending ones
+        ckpt_files = checkpoint_filesystem.get_file_info(
+            FileSelector(
+                checkpoint_path_unwrapped, recursive=False, allow_not_found=True
+            )
+        )
+        pending_suffix = f"{PENDING_CHECKPOINT_SUFFIX}.parquet"
+        pending_file_paths = [
+            f
+            for f in ckpt_files
+            if f.type == FileType.File and f.path.endswith(pending_suffix)
+        ]
+
+        if not pending_file_paths:
+            return 0
+
+        # 2. Build prefix trie from pending checkpoint basenames
+        trie = build_pending_checkpoint_trie(pending_file_paths, pending_suffix)
+
+        # 3. List all data files (recursively for partitions)
+        data_files = data_file_filesystem.get_file_info(
+            FileSelector(data_file_dir_unwrapped, recursive=True, allow_not_found=True)
+        )
+
+        # 4. Delete data files matching a pending checkpoint prefix
+        for f in data_files:
+            if f.type != FileType.File:
+                continue
+            basename = posixpath.basename(f.path)
+            if trie.has_prefix_of(basename):
+                data_file_filesystem.delete_file(f.path)
+
+        # 5. Delete all pending checkpoint files
+        for f in pending_file_paths:
+            checkpoint_filesystem.delete_file(f.path)
+
+        return len(pending_file_paths)
+
+    return call_with_retry(
+        _clean,
+        description="clean pending checkpoints",
+        max_attempts=CHECKPOINT_RECOVERY_MAX_ATTEMPTS,
+        max_backoff_s=CHECKPOINT_RECOVERY_MAX_BACKOFF_S,
+    )
+
+
+@ray.remote(num_returns=2)
+def convert_and_sort_checkpointed_ids(
+    checkpointed_ids_arrow: Block, id_column: str
+) -> Tuple[np.ndarray, int]:
+    """Convert checkpointed IDs from pyarrow.Table to sorted np.ndarray.
+
+    Args:
+        checkpointed_ids_arrow: A pyarrow.Table containing the checkpointed
+            IDs, loaded from the checkpoint parquet files.
+        id_column: The id column of `checkpoint_ids_array`.
+
+    Returns:
+        A tuple of:
+        - The sorted checkpointed IDs of type numpy.ndarray, which can be
+          passed directly to each checkpoint filter actor.
+        - The size (bytes) of the ndarray, which can be used to determine
+          the `ray_remote_args` of each checkpoint filter actor.
+    """
+    checkpointed_ids_ndarray = np.array([])
+
+    try:
+        if checkpointed_ids_arrow.num_rows != 0:
+            checkpointed_ids_ndarray = np.sort(
+                transform_pyarrow.to_numpy(
+                    checkpointed_ids_arrow[id_column], zero_copy_only=False
+                )
+            )
+    except Exception as e:
+        raise RuntimeError(f"Failed to convert and sort checkpointed IDs: {e}")
+
+    checkpoint_size = _numpy_size(checkpointed_ids_ndarray)
+    return checkpointed_ids_ndarray, checkpoint_size
+
+
+class CheckpointManager(abc.ABC):
+    """Manage checkpoint data."""
+
+    def __init__(
+        self,
+        checkpoint_config: CheckpointConfig,
+        data_context: DataContext,
+    ):
+        """Initialize the CheckpointManager.
+
+        Args:
+            checkpoint_config: the checkpoint config.
+            data_context: the DataContext snapshot whose ``execution_options``
+                should govern the Ray tasks fired during checkpoint loading
+                and pending-checkpoint cleanup. Pass the dataset's
+                ``_context`` (not ``DataContext.get_current()``) so the
+                label_selector and other execution options stay consistent
+                with the rest of materialize.
+        """
+        self.checkpoint_path = checkpoint_config.checkpoint_path
+        self.filesystem = checkpoint_config.filesystem
+        self.id_column = checkpoint_config.id_column
+        self.checkpoint_path_partition_filter = (
+            checkpoint_config.checkpoint_path_partition_filter
+        )
+        self.checkpoint_path_unwrapped = _unwrap_protocol(
+            checkpoint_config.checkpoint_path
+        )
+        self._data_context = data_context
+
+    def load_checkpoint(
+        self,
+        data_file_dir: Optional[str] = None,
+        data_file_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ) -> Tuple[Optional[ObjectRef[np.ndarray]], int]:
+        """Loading checkpoint data.
+
+        This method first cleans up any pending checkpoints from incomplete
+        2-phase commits, then loads the committed checkpoint data.
+
+        Args:
+            data_file_dir: Optional directory where data files are written.
+                If provided, pending checkpoints will be used to find and
+                delete matching data files before loading.
+            data_file_filesystem: Optional filesystem for data files. If not
+                provided, defaults to the checkpoint filesystem. Should be
+                provided when data files are on a different filesystem than
+                checkpoints.
+
+        Returns:
+            ObjectRef: The ref of checkpointed IDs array. None if no checkpoint was loaded.
+            int: the size of the checkpointed IDs array.
+        """
+        logger.info(
+            "Loading checkpoint from %s, this could take a while.", self.checkpoint_path
+        )
+
+        start_t = time.time()
+
+        # Clean up pending checkpoints before loading (runs as a Ray task)
+        if data_file_dir is not None:
+            self._clean_pending_checkpoints(data_file_dir, data_file_filesystem)
+
+        # If the checkpoint directory has no remaining data files (e.g., all
+        # entries were pending checkpoints that were just cleaned up), skip
+        # the inner ``read_parquet``. V2's ``read_parquet`` raises on empty
+        # directories while V1 returned a zero-row dataset; this pre-check
+        # keeps ``load_checkpoint`` behaving the same under both.
+        # Recurse when a partition filter is configured because committed
+        # files live under Hive-partitioned subdirectories rather than at
+        # the top level.
+        entries = self.filesystem.get_file_info(
+            FileSelector(
+                self.checkpoint_path_unwrapped,
+                recursive=self.checkpoint_path_partition_filter is not None,
+                allow_not_found=True,
+            )
+        )
+        if not any(f.type == FileType.File for f in entries):
+            return None, 0
+
+        # Load the checkpoint data
+        checkpoint_ds: ray.data.Dataset = ray.data.read_parquet(
+            self.checkpoint_path,
+            filesystem=self.filesystem,
+            partition_filter=self.checkpoint_path_partition_filter,
+        )
+        checkpoint_ds.set_name("checkpoint_dataset")
+
+        # Manually disable checkpointing for loading the checkpoint metadata
+        # to avoid recursively restoring checkpoints.
+        # TODO: Clean way to do this would be to introduce per Op config
+        # [https://github.com/ray-project/ray/issues/54520]
+        checkpoint_ds.context.checkpoint_config = None
+
+        # Pre-process data pipeline
+        checkpoint_ds: ray.data.Dataset = self._preprocess_data_pipeline(checkpoint_ds)
+
+        # Repartition to 1 block.
+        checkpoint_ds = checkpoint_ds.repartition(num_blocks=1)
+
+        # Get the block reference
+        ref_bundles: List[RefBundle] = list(checkpoint_ds.iter_internal_ref_bundles())
+
+        assert len(ref_bundles) == 1
+
+        # If there are no valid files under the checkpoint_path, return None, 0.
+        if ref_bundles[0].num_rows() == 0:
+            return None, 0
+
+        ref_bundle: RefBundle = ref_bundles[0]
+        schema: Schema = ref_bundle.schema
+        assert len(ref_bundle.blocks) == 1
+        block_ref: ObjectRef[Block] = ref_bundle.blocks[0].ref
+        metadata: BlockMetadata = ref_bundle.blocks[0].metadata
+        # Validate the loaded checkpoint
+        self._validate_loaded_checkpoint(schema, metadata)
+
+        # Convert arrow-typed ids to sorted numpy-typed ids.
+        # Note: the convert is very time-consuming.
+        # Get the object ref the checkpointed IDs, because we do not want the IDs
+        # to occupy the memory of the head node.
+        ctx_label_selector = self._data_context.execution_options.label_selector
+        task = convert_and_sort_checkpointed_ids
+        if ctx_label_selector:
+            task = task.options(label_selector=ctx_label_selector)
+        (
+            checkpointed_ids_ref,
+            checkpoint_size_ref,
+        ) = task.remote(block_ref, self.id_column)
+
+        checkpoint_size = ray.get(checkpoint_size_ref)
+
+        logger.info(
+            "Checkpoint loaded for %s in %.2f seconds. SizeBytes = %d, Schema = %s",
+            type(self).__name__,
+            time.time() - start_t,
+            checkpoint_size,
+            schema.to_string(),
+        )
+        return checkpointed_ids_ref, checkpoint_size
+
+    def _clean_pending_checkpoints(
+        self,
+        data_file_dir: Optional[str],
+        data_file_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ) -> None:
+        """Clean up pending checkpoints from incomplete 2-phase commits.
+
+        Finds pending checkpoint files, builds a prefix trie from their basenames,
+        deletes matching data files, then deletes the pending checkpoints.
+
+        Runs as a Ray task to avoid blocking the driver during potentially
+        slow filesystem operations (especially on cloud storage like S3).
+
+        Args:
+            data_file_dir: The directory where data files are written.
+            data_file_filesystem: The filesystem for data files. If not
+                provided, defaults to the checkpoint filesystem.
+        """
+        if not data_file_dir:
+            return
+        if data_file_filesystem is None:
+            data_file_filesystem = self.filesystem
+        ctx_label_selector = self._data_context.execution_options.label_selector
+        task = _clean_pending_checkpoints_task
+        if ctx_label_selector:
+            task = task.options(label_selector=ctx_label_selector)
+        try:
+            cleaned_count = ray.get(
+                task.remote(
+                    self.checkpoint_path_unwrapped,
+                    self.filesystem,
+                    _unwrap_protocol(data_file_dir),
+                    data_file_filesystem,
+                )
+            )
+            if cleaned_count > 0:
+                logger.info(f"Cleaned up {cleaned_count} pending checkpoint(s)")
+        except ray.exceptions.RayTaskError:
+            logger.exception("Failed to clean up pending checkpoints")
+            raise
+
+    def _preprocess_data_pipeline(
+        self, checkpoint_ds: ray.data.Dataset
+    ) -> ray.data.Dataset:
+        """Pre-process the checkpoint dataset.
+
+        Subclasses can override this method for custom processing.
+        """
+        return checkpoint_ds
+
+    def _validate_loaded_checkpoint(
+        self, schema: Schema, metadata: BlockMetadata
+    ) -> None:
+        """Validate the loaded checkpoint. Subclasses can override for custom validation."""
+        pass
+
+
+class IdColumnCheckpointManager(CheckpointManager):
+    """Manager for regular ID columns."""
+
+
+class CheckpointFilter(abc.ABC):
+    """Abstract class which defines the interface for filtering checkpointed rows
+    based on varying backends.
+    """
+
+    def __init__(self, config: CheckpointConfig):
+        self.ckpt_config = config
+        self.id_column = self.ckpt_config.id_column
+
+    @abstractmethod
+    def filter_rows_for_block(self, block: Block) -> Block:
+        """For the given block, filter out rows that have already
+        been checkpointed, and return the resulting block.
+
+        Args:
+            block: The input block to filter.
+        Returns:
+            A new block with rows that have not been checkpointed.
+        """
+        raise NotImplementedError
+
+
+class NumpyArrayBasedCheckpointFilter(CheckpointFilter):
+    """CheckpointFilter for batch-based backends.
+
+    This filter will first fetch the checkpointed IDs (as NumPy arrays) from the object store.
+    For each input block, it filters the block and returns the filtered block.
+    """
+
+    def __init__(
+        self,
+        checkpoint_config: CheckpointConfig,
+        checkpoint_ref: ObjectRef[np.ndarray],
+    ):
+        super().__init__(checkpoint_config)
+        self.checkpointed_ids = ray.get(checkpoint_ref)
+        assert isinstance(self.checkpointed_ids, np.ndarray)
+
+    def filter_rows_for_block(
+        self,
+        block: Block,
+    ) -> Block:
+        """Filter IDs in memory using NumPy's binary search."""
+
+        if self.checkpointed_ids.shape[0] == 0 or len(block) == 0:
+            return block
+
+        assert isinstance(block, pyarrow.Table)
+
+        # The checkpointed_ids block is sorted (see load_checkpoint).
+        # We'll use binary search to filter out processed rows.
+
+        # Convert the block's ID column to a numpy array for fast processing.
+        block_ids = transform_pyarrow.to_numpy(
+            block[self.id_column], zero_copy_only=False
+        )
+
+        # Start with a mask of all True (keep all rows).
+        mask = np.ones(len(block_ids), dtype=bool)
+        # Use binary search to find where block_ids would be in ckpt_ids.
+        sorted_indices = np.searchsorted(self.checkpointed_ids, block_ids)
+        # Only consider indices that are within bounds.
+        valid_indices = sorted_indices < len(self.checkpointed_ids)
+        # For valid indices, check for exact matches.
+        potential_matches = sorted_indices[valid_indices]
+        matched = self.checkpointed_ids[potential_matches] == block_ids[valid_indices]
+        # Mark matched IDs as False (filter out these rows).
+        mask[valid_indices] = ~matched
+
+        # Convert the final mask to a PyArrow array and filter the block.
+        mask_array = pyarrow.array(mask)
+        filtered_block = block.filter(mask_array)
+        return filtered_block

python/ray/data/checkpoint/checkpoint_writer.py

@@ -0,0 +1,273 @@
+import logging
+import os
+import uuid
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+from pyarrow import parquet as pq
+from pyarrow.fs import FileType
+
+if TYPE_CHECKING:
+    import pyarrow
+
+from ray.data._internal.util import call_with_retry
+from ray.data.block import BlockAccessor
+from ray.data.checkpoint import CheckpointBackend, CheckpointConfig
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import _unwrap_protocol
+
+logger = logging.getLogger(__name__)
+
+# Suffix for pending checkpoint files (2-phase commit)
+PENDING_CHECKPOINT_SUFFIX = ".pending"
+
+
+@dataclass
+class PendingCheckpoint:
+    """Represents a pending checkpoint file for 2-phase commit.
+
+    Attributes:
+        pending_path: Path to the pending checkpoint file.
+        committed_path: Path where the checkpoint will be after commit.
+    """
+
+    pending_path: str
+    committed_path: str
+
+
+class CheckpointWriter:
+    """Abstract class which defines the interface for writing row-level
+    checkpoints based on varying backends.
+
+    Subclasses must implement `.write_block_checkpoint()`.
+
+    For 2-phase commit support, subclasses should also implement:
+    - `.write_pending_checkpoint()`: Write checkpoint as pending file
+    - `.commit_checkpoint()`: Rename pending to committed
+    """
+
+    def __init__(self, config: CheckpointConfig):
+        self.ckpt_config = config
+        self.checkpoint_path_unwrapped = _unwrap_protocol(
+            self.ckpt_config.checkpoint_path
+        )
+        self.id_col = self.ckpt_config.id_column
+        self.filesystem = self.ckpt_config.filesystem
+        self.write_num_threads = self.ckpt_config.write_num_threads
+
+    @abstractmethod
+    def write_block_checkpoint(self, block: BlockAccessor):
+        """Write a checkpoint for all rows in a single block to the checkpoint
+        output directory given by `self.checkpoint_path`.
+
+        This is used for non-file datasinks (SQL, MongoDB, etc.) where there's
+        no predictable file path to store in checkpoint metadata. For file-based
+        datasinks that need 2-phase commit with data file path tracking, use
+        `write_pending_checkpoint()` and `commit_checkpoint()` instead.
+
+        Args:
+            block: The block accessor containing the data to checkpoint.
+
+        Subclasses of `CheckpointWriter` must implement this method."""
+        ...
+
+    def write_pending_checkpoint(
+        self,
+        id_column_data: "pyarrow.Array",
+        checkpoint_id: str,
+    ) -> Optional[PendingCheckpoint]:
+        """Write a pending checkpoint for 2-phase commit.
+
+        This is called BEFORE the data file is written. The checkpoint filename
+        is deterministic (based on checkpoint_id), enabling idempotent writes
+        on retry. For file-based datasinks, the checkpoint filename matches
+        the data file prefix, enabling recovery to match pending checkpoints
+        to data files via prefix trie.
+
+        Args:
+            id_column_data: PyArrow array containing the ID column values.
+            checkpoint_id: Deterministic identifier for the checkpoint file,
+                derived from write_uuid and task_idx. Must be the same on retry
+                to ensure idempotent writes.
+
+        Returns:
+            PendingCheckpoint object for later commit, or None if empty.
+        """
+        raise NotImplementedError(
+            "2-phase commit not implemented for this checkpoint writer"
+        )
+
+    def commit_checkpoint(self, pending: PendingCheckpoint) -> None:
+        """Commit a pending checkpoint by renaming it to committed.
+
+        This is called AFTER the data file is successfully written.
+
+        Args:
+            pending: The PendingCheckpoint to commit.
+        """
+        raise NotImplementedError(
+            "2-phase commit not implemented for this checkpoint writer"
+        )
+
+    @staticmethod
+    def create(config: CheckpointConfig) -> "CheckpointWriter":
+        """Factory method to create a `CheckpointWriter` based on the
+        provided `CheckpointConfig`."""
+        backend = config.backend
+
+        if backend in [
+            CheckpointBackend.CLOUD_OBJECT_STORAGE,
+            CheckpointBackend.FILE_STORAGE,
+        ]:
+            return BatchBasedCheckpointWriter(config)
+        raise NotImplementedError(f"Backend {backend} not implemented")
+
+
+class BatchBasedCheckpointWriter(CheckpointWriter):
+    """CheckpointWriter for batch-based backends."""
+
+    def __init__(self, config: CheckpointConfig):
+        super().__init__(config)
+
+        self.filesystem.create_dir(self.checkpoint_path_unwrapped, recursive=True)
+
+    def _prepare_checkpoint_table_from_block(self, block: BlockAccessor):
+        """Prepare the checkpoint table from a block.
+
+        Args:
+            block: The block accessor containing the data to checkpoint.
+
+        Returns:
+            PyArrow table with checkpoint IDs.
+        """
+        checkpoint_ids_block = block.select(columns=[self.id_col])
+        # `pyarrow.parquet.write_parquet` requires a PyArrow table.
+        return BlockAccessor.for_block(checkpoint_ids_block).to_arrow()
+
+    def _prepare_checkpoint_table_from_id_column(self, id_column_data: "pyarrow.Array"):
+        """Prepare the checkpoint table from ID column data.
+
+        Args:
+            id_column_data: PyArrow array containing the ID column values.
+
+        Returns:
+            PyArrow table with checkpoint IDs.
+        """
+        import pyarrow as pa
+
+        return pa.table({self.id_col: id_column_data})
+
+    def write_block_checkpoint(self, block: BlockAccessor) -> None:
+        """Write a checkpoint for all rows in a single block to the checkpoint
+        output directory given by `self.checkpoint_path`.
+
+        This is used for non-file datasinks (SQL, MongoDB, etc.) where there's
+        no predictable file path to store in checkpoint metadata.
+
+        Args:
+            block: The block accessor containing the data to checkpoint.
+        """
+        if block.num_rows() == 0:
+            return
+
+        file_name = f"{uuid.uuid4()}.parquet"
+        ckpt_file_path = os.path.join(self.checkpoint_path_unwrapped, file_name)
+
+        checkpoint_ids_table = self._prepare_checkpoint_table_from_block(block)
+
+        def _write():
+            pq.write_table(
+                checkpoint_ids_table,
+                ckpt_file_path,
+                filesystem=self.filesystem,
+            )
+
+        try:
+            call_with_retry(
+                _write,
+                description=f"Write checkpoint file: {file_name}",
+                match=DataContext.get_current().retried_io_errors,
+            )
+        except Exception:
+            logger.exception(f"Checkpoint write failed: {file_name}")
+            raise
+
+    def write_pending_checkpoint(
+        self,
+        id_column_data,
+        checkpoint_id: str,
+    ) -> Optional[PendingCheckpoint]:
+        if len(id_column_data) == 0:
+            return None
+        pending_file_name = f"{checkpoint_id}{PENDING_CHECKPOINT_SUFFIX}.parquet"
+        committed_file_name = f"{checkpoint_id}.parquet"
+
+        pending_path = os.path.join(self.checkpoint_path_unwrapped, pending_file_name)
+        committed_path = os.path.join(
+            self.checkpoint_path_unwrapped, committed_file_name
+        )
+
+        checkpoint_ids_table = self._prepare_checkpoint_table_from_id_column(
+            id_column_data
+        )
+
+        def _write():
+            pq.write_table(
+                checkpoint_ids_table,
+                pending_path,
+                filesystem=self.filesystem,
+            )
+
+        call_with_retry(
+            _write,
+            description=f"Write pending checkpoint file: {pending_file_name}",
+            match=DataContext.get_current().retried_io_errors,
+        )
+        return PendingCheckpoint(
+            pending_path=pending_path,
+            committed_path=committed_path,
+        )
+
+    def commit_checkpoint(self, pending: PendingCheckpoint) -> None:
+        """Commit a pending checkpoint by renaming it to committed.
+
+        This is called AFTER the data file is successfully written.
+
+        This operation is idempotent: if the committed file already exists
+        (and pending doesn't), it's considered already committed. This handles
+        the case where a retry happens after successful commit (e.g., network
+        timeout after move succeeded but before acknowledgment).
+
+        Args:
+            pending: The PendingCheckpoint to commit.
+        """
+
+        def _rename():
+            # Check if already committed (idempotent)
+            committed_info = self.filesystem.get_file_info(pending.committed_path)
+            pending_info = self.filesystem.get_file_info(pending.pending_path)
+
+            committed_exists = committed_info.type != FileType.NotFound
+            pending_exists = pending_info.type != FileType.NotFound
+
+            if committed_exists:
+                # Already committed. Clean up pending file if it exists.
+                if pending_exists:
+                    self.filesystem.delete_file(pending.pending_path)
+                return
+
+            if not pending_exists:
+                raise FileNotFoundError(
+                    f"Neither pending ({pending.pending_path}) nor committed "
+                    f"({pending.committed_path}) checkpoint exists"
+                )
+
+            # Normal case: move pending to committed
+            self.filesystem.move(pending.pending_path, pending.committed_path)
+
+        call_with_retry(
+            _rename,
+            description=f"Commit checkpoint: {pending.pending_path}",
+            match=DataContext.get_current().retried_io_errors,
+        )

python/ray/data/checkpoint/interfaces.py

@@ -0,0 +1,177 @@
+import os
+import warnings
+from enum import Enum
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import pyarrow
+
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.datasource import PathPartitionFilter
+
+
+@PublicAPI(stability="alpha")
+class CheckpointBackend(Enum):
+    """Supported backends for storing and reading checkpoint files.
+
+    Currently, only one type of backend is supported:
+
+    * Batch-based backends: CLOUD_OBJECT_STORAGE and FILE_STORAGE.
+
+    Their differences are as follows:
+
+    1. Writing checkpoints: Batch-based backends write a checkpoint file
+       for each block.
+    2. Loading checkpoints and filtering input data: Batch-based backends
+       load all checkpoint data into memory prior to dataset execution.
+       The checkpoint data is then passed to each read task to perform filtering.
+    """
+
+    CLOUD_OBJECT_STORAGE = "CLOUD_OBJECT_STORAGE"
+    """
+    Batch-based checkpoint backend that uses cloud object storage, such as
+    AWS S3, Google Cloud Storage, etc.
+    """
+
+    FILE_STORAGE = "FILE_STORAGE"
+    """
+    Batch based checkpoint backend that uses file system storage.
+    Note, when using this backend, the checkpoint path must be a network-mounted
+    file system (e.g. `/mnt/cluster_storage/`).
+    """
+
+
+@PublicAPI(stability="beta")
+class CheckpointConfig:
+    """Configuration for checkpointing.
+
+    Args:
+        id_column: Name of the ID column in the input dataset.
+            ID values must be unique across all rows in the dataset and must persist
+            during all operators.
+        checkpoint_path: Path to store the checkpoint data. It can be a path to a cloud
+            object storage (e.g. `s3://bucket/path`) or a file system path.
+            If the latter, the path must be a network-mounted file system (e.g.
+            `/mnt/cluster_storage/`) that is accessible to the entire cluster.
+            If not set, defaults to `RAY_DATA_CHECKPOINT_PATH_BUCKET/ray_data_checkpoint`.
+        delete_checkpoint_on_success: If true, automatically delete checkpoint
+            data when the dataset execution succeeds. Only supported for
+            batch-based backend currently.
+        override_filesystem: Override the :class:`pyarrow.fs.FileSystem` object used to
+            read/write checkpoint data. Use this when you want to use custom credentials.
+        override_backend: Override the :class:`CheckpointBackend` object used to
+            access the checkpoint backend storage.
+        write_num_threads: Number of threads used to write checkpoint files for
+            completed rows.
+        checkpoint_path_partition_filter: Filter for checkpoint files to load during
+            restoration when reading from `checkpoint_path`.
+    """
+
+    DEFAULT_CHECKPOINT_PATH_BUCKET_ENV_VAR = "RAY_DATA_CHECKPOINT_PATH_BUCKET"
+    DEFAULT_CHECKPOINT_PATH_DIR = "ray_data_checkpoint"
+    CHECKPOINT_ACTOR_POOL_MIN_SIZE = 1
+    CHECKPOINT_ACTOR_POOL_MAX_SIZE = 10
+    CHECKPOINT_ACTOR_MEMORY_BYTES = 1 * 1024**3
+
+    def __init__(
+        self,
+        id_column: Optional[str] = None,
+        checkpoint_path: Optional[str] = None,
+        *,
+        delete_checkpoint_on_success: bool = True,
+        override_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        override_backend: Optional[CheckpointBackend] = None,
+        write_num_threads: int = 3,
+        checkpoint_path_partition_filter: Optional["PathPartitionFilter"] = None,
+    ):
+        self.id_column: Optional[str] = id_column
+
+        if not isinstance(self.id_column, str) or len(self.id_column) == 0:
+            raise InvalidCheckpointingConfig(
+                "Checkpoint ID column must be a non-empty string, "
+                f"but got {self.id_column}"
+            )
+
+        if override_backend is not None:
+            warnings.warn(
+                "`override_backend` is deprecated and will be removed in August 2025.",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        self.checkpoint_path: str = (
+            checkpoint_path or self._get_default_checkpoint_path()
+        )
+        inferred_backend, inferred_fs = self._infer_backend_and_fs(
+            self.checkpoint_path,
+            override_filesystem,
+            override_backend,
+        )
+        self.filesystem: "pyarrow.fs.FileSystem" = inferred_fs
+        self.backend: CheckpointBackend = inferred_backend
+        self.delete_checkpoint_on_success: bool = delete_checkpoint_on_success
+        self.write_num_threads: int = write_num_threads
+        self.checkpoint_path_partition_filter = checkpoint_path_partition_filter
+        self.checkpoint_actor_pool_min_size = self.CHECKPOINT_ACTOR_POOL_MIN_SIZE
+        self.checkpoint_actor_pool_max_size = self.CHECKPOINT_ACTOR_POOL_MAX_SIZE
+        self.checkpoint_actor_memory_bytes = self.CHECKPOINT_ACTOR_MEMORY_BYTES
+
+    def _get_default_checkpoint_path(self) -> str:
+        artifact_storage = os.environ.get(self.DEFAULT_CHECKPOINT_PATH_BUCKET_ENV_VAR)
+        if artifact_storage is None:
+            raise InvalidCheckpointingConfig(
+                f"`{self.DEFAULT_CHECKPOINT_PATH_BUCKET_ENV_VAR}` env var is not set, "
+                "please explicitly set `CheckpointConfig.checkpoint_path`."
+            )
+        return f"{artifact_storage}/{self.DEFAULT_CHECKPOINT_PATH_DIR}"
+
+    def _infer_backend_and_fs(
+        self,
+        checkpoint_path: str,
+        override_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        override_backend: Optional[CheckpointBackend] = None,
+    ) -> Tuple[CheckpointBackend, "pyarrow.fs.FileSystem"]:
+        try:
+            if override_filesystem is not None:
+                assert isinstance(override_filesystem, pyarrow.fs.FileSystem), (
+                    "override_filesystem must be an instance of "
+                    f"`pyarrow.fs.FileSystem`, but got {type(override_filesystem)}"
+                )
+                fs = override_filesystem
+            else:
+                fs, _ = pyarrow.fs.FileSystem.from_uri(checkpoint_path)
+
+            if override_backend is not None:
+                assert isinstance(override_backend, CheckpointBackend), (
+                    "override_backend must be an instance of `CheckpointBackend`, "
+                    f"but got {type(override_backend)}"
+                )
+                backend = override_backend
+            else:
+                if isinstance(fs, pyarrow.fs.LocalFileSystem):
+                    backend = CheckpointBackend.FILE_STORAGE
+                else:
+                    backend = CheckpointBackend.CLOUD_OBJECT_STORAGE
+
+            return backend, fs
+        except Exception as e:
+            raise InvalidCheckpointingConfig(
+                f"Invalid checkpoint path: {checkpoint_path}. "
+            ) from e
+
+
+@DeveloperAPI
+class InvalidCheckpointingConfig(Exception):
+    """Exception which indicates that the checkpointing
+    configuration is invalid."""
+
+    pass
+
+
+@DeveloperAPI
+class InvalidCheckpointingOperators(Exception):
+    """Exception which indicates that the DAG is not eligible for checkpointing,
+    due to one or more incompatible operators."""
+
+    pass

python/ray/data/checkpoint/load_checkpoint_callback.py

@@ -0,0 +1,47 @@
+import logging
+
+from ray.data._internal.execution.execution_callback import (
+    ExecutionCallback,
+)
+from ray.data._internal.execution.streaming_executor import StreamingExecutor
+from ray.data.checkpoint import CheckpointConfig
+from ray.data.datasource.path_util import _unwrap_protocol
+
+logger = logging.getLogger(__name__)
+
+
+class LoadCheckpointCallback(ExecutionCallback):
+    """
+    ExecutionCallback that handles checkpoints. This Callback is responsible for
+    deleting the checkpoint directory when these conditions are met:
+    1. `delete_checkpoint_on_success` is True.
+    2. The job finishes successfully.
+    """
+
+    def __init__(
+        self,
+        config: CheckpointConfig,
+    ):
+        assert config is not None
+        self._config = config
+
+    def before_execution_starts(self, executor: StreamingExecutor):
+        assert self._config is executor._data_context.checkpoint_config
+
+    def _delete_checkpoint(self):
+        checkpoint_path_unwrapped = _unwrap_protocol(self._config.checkpoint_path)
+        filesystem = self._config.filesystem
+        filesystem.delete_dir(checkpoint_path_unwrapped)
+
+    def after_execution_succeeds(self, executor: StreamingExecutor):
+        assert self._config is executor._data_context.checkpoint_config
+
+        # Delete checkpoint data.
+        try:
+            if self._config.delete_checkpoint_on_success:
+                self._delete_checkpoint()
+        except Exception:
+            logger.warning("Failed to delete checkpoint data.", exc_info=True)
+
+    def after_execution_fails(self, executor: StreamingExecutor, error: Exception):
+        assert self._config is executor._data_context.checkpoint_config

python/ray/data/checkpoint/util.py

@@ -0,0 +1,46 @@
+import logging
+import posixpath
+from typing import List
+
+logger = logging.getLogger(__name__)
+
+
+class PrefixTrie:
+    """Trie for efficient prefix matching of filenames during recovery."""
+
+    def __init__(self) -> None:
+        self.children: dict[str, "PrefixTrie"] = {}
+        self.is_end: bool = False
+
+    def insert(self, word: str) -> None:
+        node = self
+        for ch in word:
+            if ch not in node.children:
+                node.children[ch] = PrefixTrie()
+            node = node.children[ch]
+        node.is_end = True
+
+    def has_prefix_of(self, word: str) -> bool:
+        """Return True if any inserted word is a prefix of `word`."""
+        node = self
+        for ch in word:
+            if node.is_end:
+                return True
+            if ch not in node.children:
+                return False
+            node = node.children[ch]
+        return node.is_end
+
+
+def build_pending_checkpoint_trie(file_paths: List, pending_suffix: str) -> PrefixTrie:
+    """Build a PrefixTrie from pending checkpoint file paths.
+
+    Strips the given pending suffix to get the data file prefix.
+    """
+    trie = PrefixTrie()
+    for f in file_paths:
+        basename = posixpath.basename(f.path)
+        if basename.endswith(pending_suffix):
+            prefix = basename[: -len(pending_suffix)]
+            trie.insert(prefix)
+    return trie