[Data] Abstracting common shuffle utility (#51237)

1. Abstracted common `shuffle` util 
2. Rebased `ArrowBlockAccessor.random_shuffle` to use it
3. Cleaning up `RetryingPyFileSystem`
4. Added `RowToBlockMapTransformFn`

---------

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Author: alexeykudinkin

python/ray/_private/utils.py

@@ -51,6 +51,10 @@
 if TYPE_CHECKING:
     from ray.runtime_env import RuntimeEnv
 
+
+INT32_MAX = (2**31) - 1
+
+
 pwd = None
 if sys.platform != "win32":
     import pwd

python/ray/data/_internal/arrow_block.py

@@ -22,6 +22,7 @@
     pyarrow_table_from_pydict,
 )
 from ray.data._internal.arrow_ops import transform_polars, transform_pyarrow
+from ray.data._internal.arrow_ops.transform_pyarrow import shuffle
 from ray.data._internal.row import TableRow
 from ray.data._internal.table_block import TableBlockAccessor, TableBlockBuilder
 from ray.data._internal.util import find_partitions
@@ -215,14 +216,7 @@ def slice(self, start: int, end: int, copy: bool = False) -> "pyarrow.Table":
         return view
 
     def random_shuffle(self, random_seed: Optional[int]) -> "pyarrow.Table":
-        num_rows = self.num_rows()
-        if num_rows == 0:
-            return pyarrow.table([])
-        random = np.random.RandomState(random_seed)
-        shuffled_indices = np.arange(num_rows)
-        # Shuffle all rows in-place
-        random.shuffle(shuffled_indices)
-        return self.take(pyarrow.array(shuffled_indices))
+        return shuffle(self._table, random_seed)
 
     def schema(self) -> "pyarrow.lib.Schema":
         return self._table.schema

python/ray/data/_internal/arrow_ops/transform_pyarrow.py

@@ -1,5 +1,5 @@
 import logging
-from typing import TYPE_CHECKING, List, Union, Dict
+from typing import TYPE_CHECKING, List, Union, Dict, Optional
 
 import numpy as np
 from packaging.version import parse as parse_version
@@ -59,7 +59,7 @@ def sort(table: "pyarrow.Table", sort_key: "SortKey") -> "pyarrow.Table":
 
 def take_table(
     table: "pyarrow.Table",
-    indices: Union[List[int], "pyarrow.Array", "pyarrow.ChunkedArray"],
+    indices: Union[List[int], np.ndarray, "pyarrow.Array", "pyarrow.ChunkedArray"],
 ) -> "pyarrow.Table":
     """Select rows from the table.
 
@@ -421,6 +421,19 @@ def _align_struct_fields(
     return aligned_blocks
 
 
+def shuffle(block: "pyarrow.Table", seed: Optional[int] = None) -> "pyarrow.Table":
+    """Shuffles provided Arrow table"""
+
+    if len(block) == 0:
+        return block
+
+    indices = np.arange(block.num_rows)
+    # Shuffle indices
+    np.random.RandomState(seed).shuffle(indices)
+
+    return take_table(block, indices)
+
+
 def concat(
     blocks: List["pyarrow.Table"], *, promote_types: bool = False
 ) -> "pyarrow.Table":

python/ray/data/_internal/datasource/parquet_datasource.py

@@ -202,7 +202,8 @@ def __init__(
         self._unresolved_paths = paths
         paths, self._filesystem = _resolve_paths_and_filesystem(paths, filesystem)
         filesystem = RetryingPyFileSystem.wrap(
-            self._filesystem, context=DataContext.get_current()
+            self._filesystem,
+            retryable_errors=DataContext.get_current().retried_io_errors,
         )
 
         # HACK: PyArrow's `ParquetDataset` errors if input paths contain non-parquet

python/ray/data/_internal/execution/operators/map_transformer.py

@@ -336,6 +336,30 @@ def __eq__(self, other):
         )
 
 
+class RowToBlockMapTransformFn(MapTransformFn):
+    """A Row-to-Batch MapTransformFn."""
+
+    def __init__(
+        self, transform_fn: MapTransformCallable[Row, Block], is_udf: bool = False
+    ):
+        self._transform_fn = transform_fn
+        super().__init__(
+            MapTransformFnDataType.Row,
+            MapTransformFnDataType.Block,
+            category=MapTransformFnCategory.DataProcess,
+            is_udf=is_udf,
+        )
+
+    def __call__(self, input: Iterable[Row], ctx: TaskContext) -> Iterable[Block]:
+        yield from self._transform_fn(input, ctx)
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, RowToBlockMapTransformFn)
+            and self._transform_fn == other._transform_fn
+        )
+
+
 class BlockMapTransformFn(MapTransformFn):
     """A block-to-block MapTransformFn."""
 

python/ray/data/_internal/util.py

@@ -1143,8 +1143,8 @@ def __init__(self, handler: "RetryingPyFileSystemHandler"):
         super().__init__(handler)
 
     @property
-    def data_context(self):
-        return self.handler.data_context
+    def retryable_errors(self) -> List[str]:
+        return self.handler._retryable_errors
 
     def unwrap(self):
         return self.handler.unwrap()
@@ -1153,13 +1153,15 @@ def unwrap(self):
     def wrap(
         cls,
         fs: "pyarrow.fs.FileSystem",
-        context: DataContext,
+        retryable_errors: List[str],
         max_attempts: int = 10,
         max_backoff_s: int = 32,
     ):
         if isinstance(fs, RetryingPyFileSystem):
             return fs
-        handler = RetryingPyFileSystemHandler(fs, context, max_attempts, max_backoff_s)
+        handler = RetryingPyFileSystemHandler(
+            fs, retryable_errors, max_attempts, max_backoff_s
+        )
         return cls(handler)
 
     def __reduce__(self):
@@ -1182,7 +1184,7 @@ class RetryingPyFileSystemHandler(pyarrow.fs.FileSystemHandler):
     def __init__(
         self,
         fs: "pyarrow.fs.FileSystem",
-        context: DataContext,
+        retryable_errors: List[str] = tuple(),
         max_attempts: int = 10,
         max_backoff_s: int = 32,
     ):
@@ -1198,20 +1200,16 @@ def __init__(
             fs, RetryingPyFileSystem
         ), "Cannot wrap a RetryingPyFileSystem"
         self._fs = fs
-        self._data_context = context
+        self._retryable_errors = retryable_errors
         self._max_attempts = max_attempts
         self._max_backoff_s = max_backoff_s
 
-    @property
-    def data_context(self):
-        return self._data_context
-
     def _retry_operation(self, operation: Callable, description: str):
         """Execute an operation with retries."""
         return call_with_retry(
             operation,
             description=description,
-            match=self._data_context.retried_io_errors,
+            match=self._retryable_errors,
             max_attempts=self._max_attempts,
             max_backoff_s=self._max_backoff_s,
         )

python/ray/data/datasource/file_based_datasource.py

@@ -142,7 +142,7 @@ def __init__(
         self._include_paths = include_paths
         paths, self._filesystem = _resolve_paths_and_filesystem(paths, filesystem)
         self._filesystem = RetryingPyFileSystem.wrap(
-            self._filesystem, context=self._data_context
+            self._filesystem, retryable_errors=self._data_context.retried_io_errors
         )
         paths, file_sizes = map(
             list,
@@ -475,28 +475,31 @@ def _wrap_s3_serialization_workaround(filesystem: "pyarrow.fs.FileSystem"):
 
     wrap_retries = False
     fs_to_be_wrapped = filesystem  # Only unwrap for S3FileSystemWrapper
-    context = None
+    retryable_errors = []
     if isinstance(fs_to_be_wrapped, RetryingPyFileSystem):
         wrap_retries = True
-        context = fs_to_be_wrapped.data_context
+        retryable_errors = fs_to_be_wrapped.retryable_errors
         fs_to_be_wrapped = fs_to_be_wrapped.unwrap()
     if isinstance(fs_to_be_wrapped, pa.fs.S3FileSystem):
         return _S3FileSystemWrapper(
-            fs_to_be_wrapped, wrap_retries=wrap_retries, context=context
+            fs_to_be_wrapped,
+            wrap_retries=wrap_retries,
+            retryable_errors=retryable_errors,
         )
     return filesystem
 
 
 def _unwrap_s3_serialization_workaround(
     filesystem: Union["pyarrow.fs.FileSystem", "_S3FileSystemWrapper"],
-    context: Optional[DataContext] = None,
 ):
     if isinstance(filesystem, _S3FileSystemWrapper):
         wrap_retries = filesystem._wrap_retries
-        context = filesystem._context
+        retryable_errors = filesystem._retryable_erros
         filesystem = filesystem.unwrap()
         if wrap_retries:
-            filesystem = RetryingPyFileSystem.wrap(filesystem, context=context)
+            filesystem = RetryingPyFileSystem.wrap(
+                filesystem, retryable_errors=retryable_errors
+            )
     return filesystem
 
 
@@ -505,11 +508,11 @@ def __init__(
         self,
         fs: "pyarrow.fs.S3FileSystem",
         wrap_retries: bool = False,
-        context: Optional[DataContext] = None,
+        retryable_errors: List[str] = tuple(),
     ):
         self._fs = fs
         self._wrap_retries = wrap_retries
-        self._context = context
+        self._retryable_erros = retryable_errors
 
     def unwrap(self):
         return self._fs
@@ -548,7 +551,9 @@ def _resolve_kwargs(
     return kwargs
 
 
-def _validate_shuffle_arg(shuffle: Optional[str]) -> None:
+def _validate_shuffle_arg(
+    shuffle: Union[Literal["files"], FileShuffleConfig, None]
+) -> None:
     if not (
         shuffle is None or shuffle == "files" or isinstance(shuffle, FileShuffleConfig)
     ):

python/ray/data/datasource/file_datasink.py

@@ -66,7 +66,7 @@ def __init__(
         self.unresolved_path = path
         paths, self.filesystem = _resolve_paths_and_filesystem(path, filesystem)
         self.filesystem = RetryingPyFileSystem.wrap(
-            self.filesystem, context=self._data_context
+            self.filesystem, retryable_errors=self._data_context.retried_io_errors
         )
         assert len(paths) == 1, len(paths)
         self.path = paths[0]

python/ray/data/tests/test_image.py

@@ -166,11 +166,12 @@ def test_random_shuffle(self, ray_start_regular_shared, restore_data_context):
             shuffle="files",
         )
 
-        # Execute 10 times to get a set of output paths.
+        # Execute 5 times to get a set of output paths.
         output_paths_list = []
-        for _ in range(10):
+        for _ in range(5):
             paths = [row["path"][-len(file_paths[0]) :] for row in ds.take_all()]
             output_paths_list.append(paths)
+
         all_paths_matched = [
             file_paths == output_paths for output_paths in output_paths_list
         ]

python/ray/data/tests/test_transform_pyarrow.py

@@ -16,6 +16,7 @@
     try_combine_chunked_columns,
     unify_schemas,
     MIN_PYARROW_VERSION_TYPE_PROMOTION,
+    shuffle,
 )
 from ray.data.block import BlockAccessor
 from ray.data.extensions import (
@@ -46,6 +47,20 @@ def test_try_defragment_table():
     assert dt == t
 
 
+def test_shuffle():
+    t = pa.Table.from_pydict(
+        {
+            "index": pa.array(list(range(10))),
+        }
+    )
+
+    shuffled = shuffle(t, seed=0xDEED)
+
+    assert shuffled == pa.Table.from_pydict(
+        {"index": pa.array([4, 3, 6, 8, 7, 1, 5, 2, 9, 0])}
+    )
+
+
 def test_arrow_concat_empty():
     # Test empty.
     assert concat([]) == []