(perf): Distribute data types inference (#1692)

* (perf): Distribute data types inference

* PR feedback: infer datatype based on first block

Author: jaidisido

awswrangler/_data_types.py

@@ -13,6 +13,7 @@
 import pyarrow.parquet
 
 from awswrangler import _utils, exceptions
+from awswrangler._distributed import engine
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -456,6 +457,7 @@ def pyarrow2pandas_extension(  # pylint: disable=too-many-branches,too-many-retu
     return None
 
 
+@engine.dispatch_on_engine
 def pyarrow_types_from_pandas(  # pylint: disable=too-many-branches,too-many-statements
     df: pd.DataFrame, index: bool, ignore_cols: Optional[List[str]] = None, index_left: bool = False
 ) -> Dict[str, pa.DataType]:

awswrangler/distributed/ray/_core.py

@@ -74,7 +74,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return wrapper
 
 
-def ray_get(futures: List[Any]) -> List[Any]:
+def ray_get(futures: Union["ray.ObjectRef[Any]", List["ray.ObjectRef[Any]"]]) -> Any:
     """
     Run ray.get on futures if distributed.
 

awswrangler/distributed/ray/_core.pyi

@@ -14,7 +14,7 @@ class RayLogger:
 
 def ray_logger(function: Callable[..., Any]) -> Callable[..., Any]: ...
 def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]: ...
-def ray_get(futures: List[Any]) -> List[Any]: ...
+def ray_get(futures: List[Any]) -> Any: ...
 def initialize_ray(
     address: Optional[str] = None,
     redis_password: Optional[str] = None,

awswrangler/distributed/ray/_register.py

@@ -1,5 +1,6 @@
 """Ray and Modin registered methods (PRIVATE)."""
 # pylint: disable=import-outside-toplevel
+from awswrangler._data_types import pyarrow_types_from_pandas
 from awswrangler._distributed import MemoryFormatEnum, engine, memory_format
 from awswrangler._utils import table_refs_to_df
 from awswrangler.distributed.ray._core import ray_remote
@@ -28,6 +29,7 @@ def register_ray() -> None:
 
     if memory_format.get() == MemoryFormatEnum.MODIN:
         from awswrangler.distributed.ray.modin._core import modin_repartition
+        from awswrangler.distributed.ray.modin._data_types import pyarrow_types_from_pandas_distributed
         from awswrangler.distributed.ray.modin._utils import _arrow_refs_to_df
         from awswrangler.distributed.ray.modin.s3._read_parquet import _read_parquet_distributed
         from awswrangler.distributed.ray.modin.s3._read_text import _read_text_distributed
@@ -39,6 +41,7 @@ def register_ray() -> None:
         from awswrangler.distributed.ray.modin.s3._write_text import _to_text_distributed
 
         for o_f, d_f in {
+            pyarrow_types_from_pandas: pyarrow_types_from_pandas_distributed,
             _read_parquet: _read_parquet_distributed,
             _read_text: _read_text_distributed,
             _to_buckets: _to_buckets_distributed,

awswrangler/distributed/ray/modin/_data_types.py

@@ -0,0 +1,25 @@
+"""Internal (private) Data Types Module."""
+from typing import Dict, List, Optional
+
+import modin.pandas as pd
+import pyarrow as pa
+import ray
+
+from awswrangler._data_types import pyarrow_types_from_pandas
+from awswrangler.distributed.ray._core import ray_get, ray_remote
+
+
+def pyarrow_types_from_pandas_distributed(
+    df: pd.DataFrame, index: bool, ignore_cols: Optional[List[str]] = None, index_left: bool = False
+) -> Dict[str, pa.DataType]:
+    """Extract the related Pyarrow data types from a pandas DataFrame."""
+    func = ray_remote(pyarrow_types_from_pandas)
+    first_block_object_ref = ray.data.from_modin(df).get_internal_block_refs()[0]
+    return ray_get(  # type: ignore
+        func(
+            df=first_block_object_ref,
+            index=index,
+            ignore_cols=ignore_cols,
+            index_left=index_left,
+        )
+    )

awswrangler/s3/_select.py

@@ -5,7 +5,7 @@
 import json
 import logging
 import pprint
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import boto3
 import pandas as pd
@@ -19,9 +19,6 @@
 from awswrangler.s3._list import _path2list
 from awswrangler.s3._read import _get_path_ignore_suffix
 
-if TYPE_CHECKING:
-    import ray  # pylint: disable=unused-import
-
 _logger: logging.Logger = logging.getLogger(__name__)
 
 _RANGE_CHUNK_SIZE: int = int(1024 * 1024)
@@ -42,7 +39,7 @@ def _select_object_content(
     boto3_session: Optional[boto3.Session],
     args: Dict[str, Any],
     scan_range: Optional[Tuple[int, int]] = None,
-) -> Union[pa.Table, "ray.ObjectRef[pa.Table]"]:
+) -> pa.Table:
     client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
 
     if scan_range:
@@ -83,7 +80,7 @@ def _select_query(
     scan_range_chunk_size: Optional[int] = None,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
-) -> List[Union[pa.Table, "ray.ObjectRef[pa.Table]"]]:
+) -> List[pa.Table]:
     bucket, key = _utils.parse_path(path)
 
     args: Dict[str, Any] = {