Repository: aws/aws-sdk-pandas

Issue:
Add `ray_modin_args` parameter for encapsulating distributed args
I'm wondering if we should extract `parallelism` and `bulk_read_parquet` into something like `ray_modin_args`. Both are only applied when we are in distributed mode, and it would make it easier to add more such arguments in the future.

_Originally posted by @LeonLuttenberger in https://github.com/aws/aws-sdk-pandas/pull/2033#discussion_r1117435858_
            


Discussion:

Resolving patch:
diff --git a/awswrangler/s3/_read_parquet.py b/awswrangler/s3/_read_parquet.py
--- a/awswrangler/s3/_read_parquet.py
+++ b/awswrangler/s3/_read_parquet.py
@@ -43,6 +43,7 @@
     _get_path_ignore_suffix,
     _get_path_root,
 )
+from awswrangler.typing import RayReadParquetSettings
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
@@ -364,8 +365,7 @@ def read_parquet(
     version_id: Optional[Union[str, Dict[str, str]]] = ...,
     chunked: Literal[False] = ...,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -391,8 +391,7 @@ def read_parquet(
     version_id: Optional[Union[str, Dict[str, str]]] = ...,
     chunked: Literal[True],
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -418,8 +417,7 @@ def read_parquet(
     version_id: Optional[Union[str, Dict[str, str]]] = ...,
     chunked: bool,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -445,8 +443,7 @@ def read_parquet(
     version_id: Optional[Union[str, Dict[str, str]]] = ...,
     chunked: int,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -473,8 +470,7 @@ def read_parquet(
     version_id: Optional[Union[str, Dict[str, str]]] = None,
     chunked: Union[bool, int] = False,
     use_threads: Union[bool, int] = True,
-    parallelism: int = -1,
-    bulk_read: bool = False,
+    ray_args: Optional[RayReadParquetSettings] = None,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = None,
@@ -563,14 +559,8 @@ def read_parquet(
         True to enable concurrent requests, False to disable multiple threads.
         If enabled, os.cpu_count() is used as the max number of threads.
         If integer is provided, specified number is used.
-    parallelism : int, optional
-        The requested parallelism of the read. Only used when `distributed` add-on is installed.
-        Parallelism may be limited by the number of files of the dataset. -1 (autodetect) by default.
-    bulk_read: bool, default False
-        True to enable a faster reading of a large number of Parquet files.
-        Offers improved performance due to not gathering the file metadata in a single node.
-        The drawback is that it does not offer schema resolution, so it should only be used when the
-        Parquet files are all uniform.
+    ray_args: typing.RayReadParquetSettings, optional
+        Params of the Ray Modin settings. Only used when distributed computing is used with Ray and Modin installed.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session is used if None is received.
     s3_additional_kwargs : Optional[Dict[str, Any]]
@@ -621,6 +611,9 @@ def read_parquet(
     >>> df = wr.s3.read_parquet(path, dataset=True, partition_filter=my_filter)
 
     """
+    ray_args = ray_args if ray_args else {}
+    bulk_read = ray_args.get("bulk_read", False)
+
     if bulk_read and validate_schema:
         exceptions.InvalidArgumentCombination("Cannot validate schema when bulk reading data files.")
 
@@ -692,7 +685,7 @@ def read_parquet(
         columns=columns,
         coerce_int96_timestamp_unit=coerce_int96_timestamp_unit,
         use_threads=use_threads,
-        parallelism=parallelism,
+        parallelism=ray_args.get("parallelism", -1),
         s3_client=s3_client,
         s3_additional_kwargs=s3_additional_kwargs,
         arrow_kwargs=arrow_kwargs,
@@ -715,8 +708,7 @@ def read_parquet_table(
     coerce_int96_timestamp_unit: Optional[str] = ...,
     chunked: Literal[False] = ...,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -738,8 +730,7 @@ def read_parquet_table(
     coerce_int96_timestamp_unit: Optional[str] = ...,
     chunked: Literal[True],
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -761,8 +752,7 @@ def read_parquet_table(
     coerce_int96_timestamp_unit: Optional[str] = ...,
     chunked: bool,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -784,8 +774,7 @@ def read_parquet_table(
     coerce_int96_timestamp_unit: Optional[str] = ...,
     chunked: int,
     use_threads: Union[bool, int] = ...,
-    parallelism: int = ...,
-    bulk_read: bool = ...,
+    ray_args: Optional[RayReadParquetSettings] = ...,
     boto3_session: Optional[boto3.Session] = ...,
     s3_additional_kwargs: Optional[Dict[str, Any]] = ...,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = ...,
@@ -809,8 +798,7 @@ def read_parquet_table(
     coerce_int96_timestamp_unit: Optional[str] = None,
     chunked: Union[bool, int] = False,
     use_threads: Union[bool, int] = True,
-    parallelism: int = -1,
-    bulk_read: bool = False,
+    ray_args: Optional[RayReadParquetSettings] = None,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
     pyarrow_additional_kwargs: Optional[Dict[str, Any]] = None,
@@ -875,15 +863,8 @@ def read_parquet_table(
         True to enable concurrent requests, False to disable multiple threads.
         If enabled, os.cpu_count() is used as the max number of threads.
         If integer is provided, specified number is used.
-    parallelism : int, optional
-        The requested parallelism of the read. Only used when `distributed` add-on is installed.
-        Parallelism may be limited by the number of files of the dataset. Auto-detect by default.
-    bulk_read: bool, default False
-        True to enable a faster reading of a large number of Parquet files.
-        Offers improved performance due to not gathering the file metadata in a single node.
-        The drawback is that it does not offer schema resolution, so it should only be used when the
-        Parquet files are all uniform.
-        Only used when `distributed` add-on is installed.
+    ray_args: typing.RayReadParquetSettings, optional
+        Params of the Ray Modin settings. Only used when distributed computing is used with Ray and Modin installed.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session is used if None is received.
     s3_additional_kwargs : Optional[Dict[str, Any]]
@@ -963,10 +944,9 @@ def read_parquet_table(
         columns=columns,
         validate_schema=validate_schema,
         coerce_int96_timestamp_unit=coerce_int96_timestamp_unit,
-        bulk_read=bulk_read,
         chunked=chunked,
         use_threads=use_threads,
-        parallelism=parallelism,
+        ray_args=ray_args,
         boto3_session=boto3_session,
         s3_additional_kwargs=s3_additional_kwargs,
         pyarrow_additional_kwargs=pyarrow_additional_kwargs,
diff --git a/awswrangler/s3/_read_text.py b/awswrangler/s3/_read_text.py
--- a/awswrangler/s3/_read_text.py
+++ b/awswrangler/s3/_read_text.py
@@ -20,6 +20,7 @@
     _union,
 )
 from awswrangler.s3._read_text_core import _read_text_file, _read_text_files_chunked
+from awswrangler.typing import RaySettings
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
@@ -82,7 +83,7 @@ def _read_text_format(
     dataset: bool,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]],
     ignore_index: bool,
-    parallelism: int,
+    ray_args: Optional[RaySettings],
     version_id: Optional[Union[str, Dict[str, str]]] = None,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
@@ -127,6 +128,7 @@ def _read_text_format(
             **args,
         )
 
+    ray_args = ray_args if ray_args else {}
     return _read_text(
         read_format,
         paths=paths,
@@ -136,7 +138,7 @@ def _read_text_format(
         s3_additional_kwargs=s3_additional_kwargs,
         dataset=dataset,
         ignore_index=ignore_index,
-        parallelism=parallelism,
+        parallelism=ray_args.get("parallelism", -1),
         version_ids=version_ids,
         pandas_kwargs=pandas_kwargs,
     )
@@ -158,7 +160,7 @@ def read_csv(
     chunksize: None = ...,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> pd.DataFrame:
     ...
@@ -180,7 +182,7 @@ def read_csv(
     chunksize: int,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Iterator[pd.DataFrame]:
     ...
@@ -202,7 +204,7 @@ def read_csv(
     chunksize: Optional[int],
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     ...
@@ -225,7 +227,7 @@ def read_csv(
     chunksize: Optional[int] = None,
     dataset: bool = False,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = None,
-    parallelism: int = -1,
+    ray_args: Optional[RaySettings] = None,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Read CSV file(s) from a received S3 prefix or list of S3 objects paths.
@@ -291,9 +293,8 @@ def read_csv(
         Ignored if `dataset=False`.
         E.g ``lambda x: True if x["year"] == "2020" and x["month"] == "1" else False``
         https://aws-sdk-pandas.readthedocs.io/en/3.0.0rc2/tutorials/023%20-%20Flexible%20Partitions%20Filter.html
-    parallelism : int, optional
-        The requested parallelism of the read. Only used when `distributed` add-on is installed.
-        Parallelism may be limited by the number of files of the dataset. Auto-detect by default.
+    ray_args: typing.RaySettings, optional
+        Params of the Ray Modin settings. Only used when distributed computing is used with Ray and Modin installed.
     pandas_kwargs :
         KEYWORD arguments forwarded to pandas.read_csv(). You can NOT pass `pandas_kwargs` explicitly, just add valid
         Pandas arguments in the function call and awswrangler will accept it.
@@ -360,7 +361,7 @@ def read_csv(
         last_modified_begin=last_modified_begin,
         last_modified_end=last_modified_end,
         ignore_index=ignore_index,
-        parallelism=parallelism,
+        ray_args=ray_args,
         **pandas_kwargs,
     )
 
@@ -380,7 +381,7 @@ def read_fwf(
     chunksize: None = ...,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> pd.DataFrame:
     ...
@@ -402,7 +403,7 @@ def read_fwf(
     chunksize: int,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Iterator[pd.DataFrame]:
     ...
@@ -424,7 +425,7 @@ def read_fwf(
     chunksize: Optional[int],
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     ...
@@ -447,7 +448,7 @@ def read_fwf(
     chunksize: Optional[int] = None,
     dataset: bool = False,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = None,
-    parallelism: int = -1,
+    ray_args: Optional[RaySettings] = None,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Read fixed-width formatted file(s) from a received S3 prefix or list of S3 objects paths.
@@ -513,9 +514,8 @@ def read_fwf(
         Ignored if `dataset=False`.
         E.g ``lambda x: True if x["year"] == "2020" and x["month"] == "1" else False``
         https://aws-sdk-pandas.readthedocs.io/en/3.0.0rc2/tutorials/023%20-%20Flexible%20Partitions%20Filter.html
-    parallelism : int, optional
-        The requested parallelism of the read. Only used when `distributed` add-on is installed.
-        Parallelism may be limited by the number of files of the dataset. Auto-detect by default.
+    ray_args: typing.RaySettings, optional
+        Params of the Ray Modin settings. Only used when distributed computing is used with Ray and Modin installed.
     pandas_kwargs:
         KEYWORD arguments forwarded to pandas.read_fwf(). You can NOT pass `pandas_kwargs` explicit, just add valid
         Pandas arguments in the function call and awswrangler will accept it.
@@ -582,7 +582,7 @@ def read_fwf(
         last_modified_end=last_modified_end,
         ignore_index=True,
         sort_index=False,
-        parallelism=parallelism,
+        ray_args=ray_args,
         **pandas_kwargs,
     )
 
@@ -603,7 +603,7 @@ def read_json(
     chunksize: None = ...,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> pd.DataFrame:
     ...
@@ -626,7 +626,7 @@ def read_json(
     chunksize: int,
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Iterator[pd.DataFrame]:
     ...
@@ -649,7 +649,7 @@ def read_json(
     chunksize: Optional[int],
     dataset: bool = ...,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = ...,
-    parallelism: int = ...,
+    ray_args: Optional[RaySettings] = ...,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     ...
@@ -673,7 +673,7 @@ def read_json(
     chunksize: Optional[int] = None,
     dataset: bool = False,
     partition_filter: Optional[Callable[[Dict[str, str]], bool]] = None,
-    parallelism: int = -1,
+    ray_args: Optional[RaySettings] = None,
     **pandas_kwargs: Any,
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Read JSON file(s) from a received S3 prefix or list of S3 objects paths.
@@ -742,9 +742,8 @@ def read_json(
         Ignored if `dataset=False`.
         E.g ``lambda x: True if x["year"] == "2020" and x["month"] == "1" else False``
         https://aws-sdk-pandas.readthedocs.io/en/3.0.0rc2/tutorials/023%20-%20Flexible%20Partitions%20Filter.html
-    parallelism : int, optional
-        The requested parallelism of the read. Only used when `distributed` add-on is installed.
-        Parallelism may be limited by the number of files of the dataset. Auto-detect by default.
+    ray_args: typing.RaySettings, optional
+        Params of the Ray Modin settings. Only used when distributed computing is used with Ray and Modin installed.
     pandas_kwargs:
         KEYWORD arguments forwarded to pandas.read_json(). You can NOT pass `pandas_kwargs` explicit, just add valid
         Pandas arguments in the function call and awswrangler will accept it.
@@ -814,6 +813,6 @@ def read_json(
         last_modified_begin=last_modified_begin,
         last_modified_end=last_modified_end,
         ignore_index=ignore_index,
-        parallelism=parallelism,
+        ray_args=ray_args,
         **pandas_kwargs,
     )
diff --git a/awswrangler/typing.py b/awswrangler/typing.py
--- a/awswrangler/typing.py
+++ b/awswrangler/typing.py
@@ -165,6 +165,29 @@ class AthenaPartitionProjectionSettings(TypedDict):
     """
 
 
+class RaySettings(TypedDict):
+    """Typed dictionary defining the settings for distributing calls using Ray."""
+
+    parallelism: NotRequired[int]
+    """
+    The requested parallelism of the read.
+    Parallelism may be limited by the number of files of the dataset.
+    Auto-detect by default.
+    """
+
+
+class RayReadParquetSettings(RaySettings):
+    """Typed dictionary defining the settings for distributing reading calls using Ray."""
+
+    bulk_read: NotRequired[bool]
+    """
+    True to enable a faster reading of a large number of Parquet files.
+    Offers improved performance due to not gathering the file metadata in a single node.
+    The drawback is that it does not offer schema resolution, so it should only be used when the
+    Parquet files are all uniform.
+    """
+
+
 class _S3WriteDataReturnValue(TypedDict):
     """Typed dictionary defining the dictionary returned by S3 write functions."""
 


Test patch:
diff --git a/tests/load/test_s3.py b/tests/load/test_s3.py
--- a/tests/load/test_s3.py
+++ b/tests/load/test_s3.py
@@ -54,7 +54,11 @@ def test_s3_read_parquet_simple(
 ) -> None:
     path = "s3://ursa-labs-taxi-data/2018/"
     with ExecutionTimer(request, data_paths=path) as timer:
-        wr.s3.read_parquet(path=path, bulk_read=bulk_read, validate_schema=validate_schema)
+        wr.s3.read_parquet(
+            path=path,
+            validate_schema=validate_schema,
+            ray_args={"bulk_read": bulk_read},
+        )
 
     assert timer.elapsed_time < benchmark_time
 
@@ -82,8 +86,8 @@ def test_s3_read_parquet_many_files(
     with ExecutionTimer(request, data_paths=paths) as timer:
         frame = wr.s3.read_parquet(
             path=paths,
-            bulk_read=bulk_read,
             validate_schema=validate_schema,
+            ray_args={"bulk_read": bulk_read},
         )
 
     num_files = len(paths)
diff --git a/tests/unit/test_s3_parquet.py b/tests/unit/test_s3_parquet.py
--- a/tests/unit/test_s3_parquet.py
+++ b/tests/unit/test_s3_parquet.py
@@ -179,7 +179,7 @@ def test_parquet_bulk_read(path: str, columns: Optional[List[str]]) -> None:
     for i in range(num_files):
         wr.s3.to_parquet(df=df, path=f"{path}{i}.parquet")
 
-    df2 = wr.s3.read_parquet(path=path, bulk_read=True, columns=columns)
+    df2 = wr.s3.read_parquet(path=path, columns=columns, ray_args={"bulk_read": True})
     assert len(df2) == num_files * len(df)
 
     expected_num_columns = len(df.columns) if columns is None else len(columns)