chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource_v2/partitioners/file_partitioner.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+
+
+class FilePartitioner(ABC):
+    """Abstract base class for partitioning file manifests.
+
+    A ``FilePartitioner`` groups file paths and their associated metadata into new
+    file manifests based on a specific partitioning strategy.
+
+    Implementations must be deterministic to ensure consistent partitioning across
+    retries.
+    """
+
+    @abstractmethod
+    def add_input(self, input_manifest: FileManifest):
+        """Add a file manifest to be partitioned.
+
+        Args:
+            input_manifest: A ``FileManifest`` containing paths and metadata to partition.
+        """
+        ...
+
+    @abstractmethod
+    def has_partition(self) -> bool:
+        """Check if there are any partitions available.
+
+        Returns:
+            ``True`` if there are partitions ready to be retrieved via
+            ``next_partition()``, ``False`` otherwise.
+        """
+        ...
+
+    @abstractmethod
+    def next_partition(self) -> FileManifest:
+        """Get the next available partition.
+
+        Returns:
+            A ``FileManifest`` containing the paths and metadata for the next partition.
+        """
+        ...
+
+    @abstractmethod
+    def finalize(self):
+        """Process any remaining files and complete the partitioning.
+
+        This method is called after all inputs have been added via ``add_input()`` to
+        ensure any buffered files are properly partitioned.
+        """
+        ...

python/ray/data/_internal/datasource_v2/partitioners/round_robin_partitioner.py

@@ -0,0 +1,89 @@
+import logging
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.partitioners.file_partitioner import (
+    FilePartitioner,
+)
+from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
+    InMemorySizeEstimator,
+)
+from ray.data._internal.weighted_round_robin import WeightedRoundRobinPartitioner
+
+logger = logging.getLogger(__name__)
+
+
+class RoundRobinPartitioner(FilePartitioner):
+    """Partitions input paths into blocks based on the in-memory size of files.
+
+    This partitioning ensures read tasks effectively utilize the cluster and
+    produce appropriately-sized blocks
+
+    **Steps:**
+        1. Initialize empty buckets.
+        2. Iterate through input blocks and add paths to buckets. For each path:
+            - If the current bucket falls below `min_bucket_size`, add the path and don't move
+              to the next bucket.
+            - If the current bucket exceeds `min_bucket_size` but not `max_bucket_size`,
+              add the path and move to the next bucket.
+            - If the current bucket exceeds `max_bucket_size`, yield the paths as a block, clear
+              the bucket, and move to the next bucket.
+        3. Yield any remaining paths in the buckets as blocks.
+
+    This algorithm ensures that each block contains [min_bucket_size, max_bucket_size]
+    worth of files.  It's a deterministic algorithm, but it doesn't maintain the order
+    of the input paths.
+    """
+
+    def __init__(
+        self,
+        in_memory_size_estimator: InMemorySizeEstimator,
+        *,
+        min_bucket_size: int,
+        max_bucket_size: int,
+        num_buckets: int,
+    ):
+        self._in_memory_size_estimator = in_memory_size_estimator
+        self._partitioner = WeightedRoundRobinPartitioner(
+            min_bucket_size=min_bucket_size,
+            max_bucket_size=max_bucket_size,
+            num_buckets=num_buckets,
+        )
+
+    def add_input(self, input_manifest: FileManifest):
+        in_memory_size_estimates = (
+            self._in_memory_size_estimator.estimate_in_memory_sizes(input_manifest)
+        )
+        for (
+            file_path,
+            file_size,
+            file_chunk_metadata,
+            in_memory_size_estimate,
+        ) in zip(
+            input_manifest.paths,
+            input_manifest.file_sizes,
+            input_manifest.file_chunk_metadatas,
+            in_memory_size_estimates,
+        ):
+            self._partitioner.add_item(
+                (file_path, file_size, file_chunk_metadata),
+                in_memory_size_estimate,
+            )
+
+    def has_partition(self) -> bool:
+        return self._partitioner.has_partition()
+
+    @property
+    def num_buckets(self) -> int:
+        return self._partitioner.num_buckets
+
+    def next_partition(self) -> FileManifest:
+        partition = self._partitioner.next_partition()
+        paths, file_sizes, file_chunk_metadatas = zip(*partition)
+        return FileManifest.construct_manifest(
+            list(paths),
+            list(file_sizes),
+            list(file_chunk_metadatas),
+        )
+
+    def finalize(self):
+        self._partitioner.finalize()