fix: replace deprecated ray parallelism arg with override_num_blocks (#2876)

Author: kukushking

awswrangler/distributed/ray/modin/s3/_read_orc.py

@@ -26,7 +26,7 @@ def _read_orc_distributed(
     schema: pa.schema | None,
     columns: list[str] | None,
     use_threads: bool | int,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     s3_client: "S3Client" | None,
     s3_additional_kwargs: dict[str, Any] | None,
@@ -43,7 +43,7 @@ def _read_orc_distributed(
     )
     ray_dataset = read_datasource(
         datasource,
-        parallelism=parallelism,
+        override_num_blocks=override_num_blocks,
     )
     to_pandas_kwargs = _data_types.pyarrow2pandas_defaults(
         use_threads=use_threads,

awswrangler/distributed/ray/modin/s3/_read_parquet.py

@@ -34,7 +34,7 @@ def _read_parquet_distributed(
     columns: list[str] | None,
     coerce_int96_timestamp_unit: str | None,
     use_threads: bool | int,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     s3_client: "S3Client" | None,
     s3_additional_kwargs: dict[str, Any] | None,
@@ -60,7 +60,7 @@ def _read_parquet_distributed(
                 "dataset_kwargs": dataset_kwargs,
             },
         ),
-        parallelism=parallelism,
+        override_num_blocks=override_num_blocks,
     )
     return _to_modin(
         dataset=dataset,

awswrangler/distributed/ray/modin/s3/_read_text.py

@@ -138,7 +138,7 @@ def _read_text_distributed(
     s3_additional_kwargs: dict[str, str] | None,
     dataset: bool,
     ignore_index: bool,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     pandas_kwargs: dict[str, Any],
 ) -> pd.DataFrame:
@@ -172,6 +172,6 @@ def _read_text_distributed(
             meta_provider=FastFileMetadataProvider(),
             **configuration,
         ),
-        parallelism=parallelism,
+        override_num_blocks=override_num_blocks,
     )
     return _to_modin(dataset=ray_dataset, ignore_index=ignore_index)

awswrangler/s3/_read.py

@@ -29,6 +29,7 @@
 from awswrangler.catalog._utils import _catalog_id
 from awswrangler.distributed.ray import ray_get
 from awswrangler.s3._list import _path2list, _prefix_cleanup
+from awswrangler.typing import RaySettings
 
 if TYPE_CHECKING:
     from mypy_boto3_glue.type_defs import GetTableResponseTypeDef
@@ -377,3 +378,20 @@ def _get_paths_for_glue_table(
                 )
 
     return paths, path_root, res
+
+
+def _get_num_output_blocks(
+    ray_args: RaySettings | None = None,
+) -> int:
+    ray_args = ray_args or {}
+    parallelism = ray_args.get("parallelism", -1)
+    override_num_blocks = ray_args.get("override_num_blocks")
+    if parallelism != -1:
+        pass
+        _logger.warning(
+            "The argument ``parallelism`` is deprecated and will be removed in the next major release. "
+            "Please specify ``override_num_blocks`` instead."
+        )
+    elif override_num_blocks is not None:
+        parallelism = override_num_blocks
+    return parallelism

awswrangler/s3/_read_orc.py

@@ -28,6 +28,7 @@
     _apply_partition_filter,
     _check_version_id,
     _extract_partitions_dtypes_from_table_details,
+    _get_num_output_blocks,
     _get_path_ignore_suffix,
     _get_path_root,
     _get_paths_for_glue_table,
@@ -137,7 +138,7 @@ def _read_orc(
     schema: pa.schema | None,
     columns: list[str] | None,
     use_threads: bool | int,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     s3_client: "S3Client" | None,
     s3_additional_kwargs: dict[str, Any] | None,
@@ -283,8 +284,6 @@ def read_orc(
     >>> df = wr.s3.read_orc(path, dataset=True, partition_filter=my_filter)
 
     """
-    ray_args = ray_args if ray_args else {}
-
     s3_client = _utils.client(service_name="s3", session=boto3_session)
     paths: list[str] = _path2list(
         path=path,
@@ -330,7 +329,7 @@ def read_orc(
         schema=schema,
         columns=columns,
         use_threads=use_threads,
-        parallelism=ray_args.get("parallelism", -1),
+        override_num_blocks=_get_num_output_blocks(ray_args),
         s3_client=s3_client,
         s3_additional_kwargs=s3_additional_kwargs,
         arrow_kwargs=arrow_kwargs,

awswrangler/s3/_read_parquet.py

@@ -34,6 +34,7 @@
     _apply_partition_filter,
     _check_version_id,
     _extract_partitions_dtypes_from_table_details,
+    _get_num_output_blocks,
     _get_path_ignore_suffix,
     _get_path_root,
     _get_paths_for_glue_table,
@@ -285,7 +286,7 @@ def _read_parquet(
     columns: list[str] | None,
     coerce_int96_timestamp_unit: str | None,
     use_threads: bool | int,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     s3_client: "S3Client" | None,
     s3_additional_kwargs: dict[str, Any] | None,
@@ -562,7 +563,7 @@ def read_parquet(
         columns=columns,
         coerce_int96_timestamp_unit=coerce_int96_timestamp_unit,
         use_threads=use_threads,
-        parallelism=ray_args.get("parallelism", -1),
+        override_num_blocks=_get_num_output_blocks(ray_args),
         s3_client=s3_client,
         s3_additional_kwargs=s3_additional_kwargs,
         arrow_kwargs=arrow_kwargs,

awswrangler/s3/_read_parquet.pyi

@@ -28,7 +28,7 @@ def _read_parquet(
     columns: list[str] | None,
     coerce_int96_timestamp_unit: str | None,
     use_threads: bool | int,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     s3_client: "S3Client" | None,
     s3_additional_kwargs: dict[str, Any] | None,

awswrangler/s3/_read_text.py

@@ -19,6 +19,7 @@
 from awswrangler.s3._read import (
     _apply_partition_filter,
     _check_version_id,
+    _get_num_output_blocks,
     _get_path_ignore_suffix,
     _get_path_root,
     _union,
@@ -52,7 +53,7 @@ def _read_text(
     s3_additional_kwargs: dict[str, str] | None,
     dataset: bool,
     ignore_index: bool,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     pandas_kwargs: dict[str, Any],
 ) -> pd.DataFrame:
@@ -131,7 +132,6 @@ def _read_text_format(
             **args,
         )
 
-    ray_args = ray_args if ray_args else {}
     return _read_text(
         read_format,
         paths=paths,
@@ -141,7 +141,7 @@ def _read_text_format(
         s3_additional_kwargs=s3_additional_kwargs,
         dataset=dataset,
         ignore_index=ignore_index,
-        parallelism=ray_args.get("parallelism", -1),
+        override_num_blocks=_get_num_output_blocks(ray_args),
         version_ids=version_ids,
         pandas_kwargs=pandas_kwargs,
     )

awswrangler/s3/_read_text.pyi

@@ -19,7 +19,7 @@ def _read_text(
     s3_additional_kwargs: dict[str, str] | None,
     dataset: bool,
     ignore_index: bool,
-    parallelism: int,
+    override_num_blocks: int,
     version_ids: dict[str, str] | None,
     pandas_kwargs: dict[str, Any],
 ) -> pd.DataFrame | Iterator[pd.DataFrame]: ...

awswrangler/typing.py

@@ -231,6 +231,13 @@ class RaySettings(TypedDict):
     Parallelism may be limited by the number of files of the dataset.
     Auto-detect by default.
     """
+    override_num_blocks: NotRequired[int]
+    """
+    Override the number of output blocks from all read tasks.
+    By default, the number of output blocks is dynamically decided based on
+    input data size and available resources. You shouldn't manually set this
+    value in most cases.
+    """
 
 
 class RayReadParquetSettings(RaySettings):

tests/glue_scripts/ray_read_small_parquet.py

@@ -5,4 +5,4 @@
 import awswrangler as wr
 
 paths = wr.s3.list_objects(f"s3://{os.environ['data-gen-bucket']}/parquet/small/partitioned/")
-ray.data.read_parquet_bulk(paths=paths, parallelism=1000).to_modin()
+ray.data.read_parquet_bulk(paths=paths, override_num_blocks=1000).to_modin()

tests/glue_scripts/wrangler_read_small_parquet.py

@@ -4,5 +4,5 @@
 
 wr.s3.read_parquet(
     path=f"s3://{os.environ['data-gen-bucket']}/parquet/small/partitioned/",
-    ray_args={"parallelism": 1000, "bulk_read": True},
+    ray_args={"override_num_blocks": 1000, "bulk_read": True},
 )

tests/glue_scripts/wrangler_write_partitioned_parquet.py

@@ -4,7 +4,7 @@
 
 df = wr.s3.read_parquet(
     path=f"s3://{os.environ['data-gen-bucket']}/parquet/medium/partitioned/",
-    ray_args={"parallelism": 1000},
+    ray_args={"override_num_blocks": 1000},
 )
 
 wr.s3.to_parquet(