chore: import upstream snapshot with attribution

python/ray/data/datasource/filename_provider.py

@@ -0,0 +1,171 @@
+from typing import Any, Dict, Optional, Tuple
+
+from ray.data.block import Block
+from ray.util.annotations import Deprecated, PublicAPI
+
+
+def _split_base_and_ext(filename: str) -> Tuple[str, str]:
+    """Split a filename into (base, extension) where extension includes the dot.
+
+    Returns (base, ext) where ext includes the leading dot (e.g., ".parquet"),
+    or is empty string if the filename has no extension.
+
+    This is the single source of truth for separating a task filename's base
+    from its extension. Used by both row-filename derivation and checkpoint
+    base-filename extraction — these MUST agree for prefix-trie recovery.
+    """
+    if "." in filename:
+        base, ext = filename.rsplit(".", 1)
+        return base, f".{ext}"
+    return filename, ""
+
+
+@PublicAPI(stability="alpha")
+class FilenameProvider:
+    """Generates filenames when you write a :class:`~ray.data.Dataset`.
+
+    Use this class to customize the filenames used when writing a Dataset.
+
+    Override :meth:`~FilenameProvider.get_filename_for_task` to customize filenames.
+    For row-based writes (e.g., :meth:`~ray.data.Dataset.write_images`), row filenames
+    are automatically derived by appending ``_{block_index:06}_{row_index:06}`` to the
+    task filename.
+
+    Example:
+
+        This snippet shows you how to customize filenames with a prefix. For example,
+        a file might be named ``images_abc123_000000.png``.
+
+        .. testcode::
+
+            import ray
+            from ray.data.datasource import FilenameProvider
+
+            class ImageFilenameProvider(FilenameProvider):
+
+                def __init__(self, prefix: str, file_format: str):
+                    super().__init__(file_format=file_format)
+                    self.prefix = prefix
+
+                def get_filename_for_task(self, write_uuid, task_index):
+                    return f"{self.prefix}_{write_uuid}_{task_index:06}.{self.file_format}"
+
+            ds = ray.data.read_parquet("s3://anonymous@ray-example-data/images.parquet")
+            ds.write_images(
+                "/tmp/results",
+                column="image",
+                filename_provider=ImageFilenameProvider("images", "png")
+            )
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        dataset_uuid: Optional[str] = None,
+        file_format: Optional[str] = None,
+    ) -> None:
+        """Create a FilenameProvider.
+
+        Args:
+            dataset_uuid: An optional UUID to include as a prefix in generated
+                filenames. If provided, filenames will be formatted as
+                ``{dataset_uuid}_{write_uuid}_{task_index}``.
+            file_format: An optional file extension (without the leading dot)
+                to append to generated filenames, e.g. ``"parquet"`` or
+                ``"csv"``.
+        """
+        self.dataset_uuid = dataset_uuid
+        self.file_format = file_format
+
+    def get_filename_for_task(self, write_uuid: str, task_index: int) -> str:
+        """Generate a filename for a write task.
+
+        Override this method to customize filenames when writing a Dataset.
+
+        .. note::
+            Filenames must be unique and deterministic for a given write UUID and
+            task index.
+
+        Args:
+            write_uuid: The UUID of the write operation.
+            task_index: The index of the write task.
+
+        Returns:
+            The generated filename string.
+        """
+        file_id = f"{write_uuid}_{task_index:06}"
+        filename = ""
+        if self.dataset_uuid is not None:
+            filename += f"{self.dataset_uuid}_"
+        filename += file_id
+        if self.file_format is not None:
+            filename += f".{self.file_format}"
+        return filename
+
+    @Deprecated(
+        message="Use get_filename_for_task() instead. The block and block_index "
+        "parameters are unused in practice because datasinks merge all blocks into "
+        "one before writing. These parameters will be removed in a future release. "
+        "Do not depend on block content or block_index in your FilenameProvider "
+        "implementation - filenames must be deterministic from (write_uuid, task_index) "
+        "alone to ensure checkpointing correctness."
+    )
+    def get_filename_for_block(
+        self, block: Optional[Block], write_uuid: str, task_index: int, block_index: int
+    ) -> str:
+        """Generate a filename for a block of data.
+
+        .. note::
+            Filenames must be unique and deterministic for a given write UUID and
+            task index. Do NOT depend on block content or block_index.
+
+            Checkpointing requires predicting the output filename BEFORE writing
+            data. This enables 2-phase commit: if a write fails after creating the
+            file but before committing the checkpoint, recovery can use the
+            predicted filename to delete orphaned files and retry cleanly. If
+            filenames depend on block content, this prediction is impossible and
+            checkpointing cannot guarantee exactly-once semantics.
+
+        Args:
+            block: Deprecated, unused. Do not depend on block content.
+            write_uuid: The UUID of the write operation.
+            task_index: The index of the write task.
+            block_index: Deprecated, always 0. Do not depend on this value.
+
+        Returns:
+            The filename to use for the block.
+        """
+        raise NotImplementedError
+
+    @Deprecated(
+        message="Implement get_filename_for_task() instead. Row filenames are "
+        "automatically derived by appending _{block_index:06}_{row_index:06} to the "
+        "task filename. All files from the same task must share the task filename as "
+        "a prefix so that uncommitted data files can be identified and cleaned up "
+        "during checkpoint recovery."
+    )
+    def get_filename_for_row(
+        self,
+        row: Dict[str, Any],
+        write_uuid: str,
+        task_index: int,
+        block_index: int,
+        row_index: int,
+    ) -> str:
+        """Generate a filename for a row.
+
+        .. deprecated::
+            Implement :meth:`get_filename_for_task` instead. Row filenames are
+            automatically derived by appending ``_{block_index:06}_{row_index:06}``
+            to the task filename.
+
+        Args:
+            row: The row that will be written to a file.
+            write_uuid: The UUID of the write operation.
+            task_index: The index of the write task.
+            block_index: The index of the block *within* the write task.
+            row_index: The index of the row *within* the block.
+
+        Returns:
+            The filename to use for the row.
+        """
+        raise NotImplementedError