chore: import upstream snapshot with attribution

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,
+        )