Optimize distributed CSV I/O by adding PyArrow-based datasource (#1699)

LeonLuttenberger · web-flow · commit 665141e59b3d · 2022-10-21T11:14:07.000-05:00
diff --git a/awswrangler/distributed/ray/datasources/__init__.py b/awswrangler/distributed/ray/datasources/__init__.py
@@ -1,5 +1,6 @@
 """Ray Datasources Module."""
 
+from awswrangler.distributed.ray.datasources.arrow_csv_datasource import ArrowCSVDatasource
 from awswrangler.distributed.ray.datasources.pandas_file_based_datasource import UserProvidedKeyBlockWritePathProvider
 from awswrangler.distributed.ray.datasources.pandas_text_datasource import (
     PandasCSVDataSource,
@@ -10,6 +11,7 @@
 from awswrangler.distributed.ray.datasources.parquet_datasource import ParquetDatasource
 
 __all__ = [
+    "ArrowCSVDatasource",
     "PandasCSVDataSource",
     "PandasFWFDataSource",
     "PandasJSONDatasource",
diff --git a/awswrangler/distributed/ray/datasources/arrow_csv_datasource.py b/awswrangler/distributed/ray/datasources/arrow_csv_datasource.py
@@ -0,0 +1,68 @@
+"""Ray ArrowCSVDatasource Module."""
+from typing import Any, Iterator
+
+import pyarrow as pa
+from pyarrow import csv
+from ray.data.block import BlockAccessor
+
+from awswrangler._arrow import _add_table_partitions
+from awswrangler.distributed.ray.datasources.pandas_file_based_datasource import PandasFileBasedDatasource
+
+
+class ArrowCSVDatasource(PandasFileBasedDatasource):  # pylint: disable=abstract-method
+    """CSV datasource, for reading and writing CSV files using PyArrow."""
+
+    _FILE_EXTENSION = "csv"
+
+    def _read_stream(  # type: ignore  # pylint: disable=arguments-differ
+        self,
+        f: pa.NativeFile,
+        path: str,
+        path_root: str,
+        dataset: bool,
+        **reader_args: Any,
+    ) -> Iterator[pa.Table]:
+        read_options = reader_args.get("read_options", csv.ReadOptions(use_threads=False))
+        parse_options = reader_args.get(
+            "parse_options",
+            csv.ParseOptions(),
+        )
+        convert_options = reader_args.get("convert_options", csv.ConvertOptions())
+
+        reader = csv.open_csv(
+            f,
+            read_options=read_options,
+            parse_options=parse_options,
+            convert_options=convert_options,
+        )
+
+        schema = None
+        while True:
+            try:
+                batch = reader.read_next_batch()
+                table = pa.Table.from_batches([batch], schema=schema)
+                if schema is None:
+                    schema = table.schema
+
+                if dataset:
+                    table = _add_table_partitions(
+                        table=table,
+                        path=f"s3://{path}",
+                        path_root=path_root,
+                    )
+
+                yield table
+
+            except StopIteration:
+                return
+
+    def _write_block(  # type: ignore  # pylint: disable=arguments-differ
+        self,
+        f: pa.NativeFile,
+        block: BlockAccessor[Any],
+        **writer_args,
+    ) -> None:
+        write_options_dict = writer_args.get("write_options", {})
+        write_options = csv.WriteOptions(**write_options_dict)
+
+        csv.write_csv(block.to_arrow(), f, write_options)
diff --git a/awswrangler/distributed/ray/modin/_utils.py b/awswrangler/distributed/ray/modin/_utils.py
@@ -1,5 +1,6 @@
 """Modin on Ray utilities (PRIVATE)."""
-from typing import Any, Callable, Dict, List, Optional, Union
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Set, Union
 
 import modin.pandas as modin_pd
 import pandas as pd
@@ -8,6 +9,7 @@
 from ray.data._internal.arrow_block import ArrowBlockAccessor, ArrowRow
 from ray.data._internal.remote_fn import cached_remote_fn
 
+from awswrangler import exceptions
 from awswrangler._arrow import _table_to_df
 
 
@@ -43,3 +45,30 @@ def _to_modin(
 
 def _arrow_refs_to_df(arrow_refs: List[Callable[..., Any]], kwargs: Optional[Dict[str, Any]]) -> modin_pd.DataFrame:
     return _to_modin(dataset=ray.data.from_arrow_refs(arrow_refs), to_pandas_kwargs=kwargs)
+
+
+@dataclass
+class ParamConfig:
+    """
+    Configuration for a Pandas argument that is supported in PyArrow.
+
+    Contains a default value and, optionally, a list of supports values.
+    """
+
+    default: Any
+    supported_values: Optional[Set[Any]] = None
+
+
+def _check_parameters(pandas_kwargs: Dict[str, Any], supported_params: Dict[str, ParamConfig]) -> None:
+    for pandas_arg_key, pandas_args_value in pandas_kwargs.items():
+        if pandas_arg_key not in supported_params:
+            raise exceptions.InvalidArgument(f"Unsupported Pandas parameter for PyArrow loader: {pandas_arg_key}")
+
+        param_config = supported_params[pandas_arg_key]
+        if param_config.supported_values is None:
+            continue
+
+        if pandas_args_value not in param_config.supported_values:
+            raise exceptions.InvalidArgument(
+                f"Unsupported Pandas parameter value for PyArrow loader: {pandas_arg_key}={pandas_args_value}",
+            )
diff --git a/awswrangler/distributed/ray/modin/s3/_read_text.py b/awswrangler/distributed/ray/modin/s3/_read_text.py
@@ -1,18 +1,70 @@
 """Modin on Ray S3 read text module (PRIVATE)."""
-from typing import Any, Dict, List, Optional, Union
+import logging
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import boto3
 import modin.pandas as pd
+from pyarrow import csv
 from ray.data import read_datasource
 
 from awswrangler import exceptions
-from awswrangler.distributed.ray.datasources import PandasCSVDataSource, PandasFWFDataSource, PandasJSONDatasource
-from awswrangler.distributed.ray.modin._utils import _to_modin
+from awswrangler.distributed.ray.datasources import (
+    ArrowCSVDatasource,
+    PandasCSVDataSource,
+    PandasFWFDataSource,
+    PandasJSONDatasource,
+)
+from awswrangler.distributed.ray.modin._utils import ParamConfig, _check_parameters, _to_modin
 
+_logger: logging.Logger = logging.getLogger(__name__)
 
-def _resolve_format(read_format: str) -> Any:
+_CSV_SUPPORTED_PARAMS = {
+    "sep": ParamConfig(default=","),
+    "delimiter": ParamConfig(default=","),
+    "quotechar": ParamConfig(default='"'),
+    "doublequote": ParamConfig(default=True),
+}
+
+
+def _parse_csv_configuration(
+    pandas_kwargs: Dict[str, Any],
+) -> Tuple[csv.ReadOptions, csv.ParseOptions, csv.ConvertOptions]:
+    _check_parameters(pandas_kwargs, _CSV_SUPPORTED_PARAMS)
+
+    read_options = csv.ReadOptions(
+        use_threads=False,
+    )
+    parse_options = csv.ParseOptions(
+        delimiter=pandas_kwargs.get("sep", _CSV_SUPPORTED_PARAMS["sep"].default),
+        quote_char=pandas_kwargs.get("quotechar", _CSV_SUPPORTED_PARAMS["quotechar"].default),
+        double_quote=pandas_kwargs.get("doublequote", _CSV_SUPPORTED_PARAMS["doublequote"].default),
+    )
+    convert_options = csv.ConvertOptions()
+
+    return read_options, parse_options, convert_options
+
+
+def _parse_configuration(
+    file_format: str,
+    version_ids: Dict[str, Optional[str]],
+    s3_additional_kwargs: Optional[Dict[str, str]],
+    pandas_kwargs: Dict[str, Any],
+) -> Tuple[csv.ReadOptions, csv.ParseOptions, csv.ConvertOptions]:
+    if {key: value for key, value in version_ids.items() if value is not None}:
+        raise exceptions.InvalidArgument("Specific version ID found for object")
+
+    if s3_additional_kwargs:
+        raise exceptions.InvalidArgument(f"Additional S3 args specified: {s3_additional_kwargs}")
+
+    if file_format == "csv":
+        return _parse_csv_configuration(pandas_kwargs)
+
+    raise exceptions.InvalidArgument()
+
+
+def _resolve_format(read_format: str, can_use_arrow: bool) -> Any:
     if read_format == "csv":
-        return PandasCSVDataSource()
+        return ArrowCSVDatasource() if can_use_arrow else PandasCSVDataSource()
     if read_format == "fwf":
         return PandasFWFDataSource()
     if read_format == "json":
@@ -33,14 +85,34 @@ def _read_text_distributed(  # pylint: disable=unused-argument
     use_threads: Union[bool, int],
     boto3_session: Optional["boto3.Session"],
 ) -> pd.DataFrame:
-    ds = read_datasource(
-        datasource=_resolve_format(read_format),
+    try:
+        read_options, parse_options, convert_options = _parse_configuration(
+            read_format,
+            version_id_dict,
+            s3_additional_kwargs,
+            pandas_kwargs,
+        )
+        can_use_arrow = True
+    except exceptions.InvalidArgument as e:
+        _logger.warning(
+            "PyArrow method unavailable, defaulting to Pandas I/O functions: %s. "
+            "This will result in slower performance of the read operations",
+            e,
+        )
+        read_options, parse_options, convert_options = None, None, None
+        can_use_arrow = False
+
+    ray_dataset = read_datasource(
+        datasource=_resolve_format(read_format, can_use_arrow),
         parallelism=parallelism,
         paths=paths,
         path_root=path_root,
         dataset=dataset,
         version_ids=version_id_dict,
         s3_additional_kwargs=s3_additional_kwargs,
         pandas_kwargs=pandas_kwargs,
+        read_options=read_options,
+        parse_options=parse_options,
+        convert_options=convert_options,
     )
-    return _to_modin(dataset=ds, ignore_index=ignore_index)
+    return _to_modin(dataset=ray_dataset, ignore_index=ignore_index)
diff --git a/awswrangler/distributed/ray/modin/s3/_write_text.py b/awswrangler/distributed/ray/modin/s3/_write_text.py
@@ -10,17 +10,63 @@
 
 from awswrangler import exceptions
 from awswrangler.distributed.ray.datasources import (  # pylint: disable=ungrouped-imports
+    ArrowCSVDatasource,
     PandasCSVDataSource,
     PandasJSONDatasource,
-    PandasTextDatasource,
     UserProvidedKeyBlockWritePathProvider,
 )
+from awswrangler.distributed.ray.datasources.pandas_file_based_datasource import PandasFileBasedDatasource
+from awswrangler.distributed.ray.modin._utils import ParamConfig, _check_parameters
 from awswrangler.s3._write import _COMPRESSION_2_EXT
 from awswrangler.s3._write_text import _get_write_details
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+_CSV_SUPPORTED_PARAMS: Dict[str, ParamConfig] = {
+    "header": ParamConfig(default=True),
+    "sep": ParamConfig(default=",", supported_values={","}),
+    "index": ParamConfig(default=True, supported_values={True}),
+    "compression": ParamConfig(default=None, supported_values={None}),
+    "quoting": ParamConfig(default=None, supported_values={None}),
+    "escapechar": ParamConfig(default=None, supported_values={None}),
+    "date_format": ParamConfig(default=None, supported_values={None}),
+}
+
+
+def _parse_csv_configuration(
+    pandas_kwargs: Dict[str, Any],
+) -> Dict[str, Any]:
+    _check_parameters(pandas_kwargs, _CSV_SUPPORTED_PARAMS)
+
+    # csv.WriteOptions cannot be pickled for some reason so we're building a Python dict
+    return {
+        "include_header": pandas_kwargs.get("header", _CSV_SUPPORTED_PARAMS["header"].default),
+    }
+
+
+def _parse_configuration(
+    file_format: str,
+    s3_additional_kwargs: Optional[Dict[str, str]],
+    pandas_kwargs: Dict[str, Any],
+) -> Dict[str, Any]:
+    if s3_additional_kwargs:
+        raise exceptions.InvalidArgument(f"Additional S3 args specified: {s3_additional_kwargs}")
+
+    if file_format == "csv":
+        return _parse_csv_configuration(pandas_kwargs)
+
+    raise exceptions.InvalidArgument()
+
+
+def _datasource_for_format(read_format: str, can_use_arrow: bool) -> PandasFileBasedDatasource:
+    if read_format == "csv":
+        return ArrowCSVDatasource() if can_use_arrow else PandasCSVDataSource()
+    if read_format == "json":
+        return PandasJSONDatasource()
+    raise exceptions.UnsupportedType("Unsupported read format")
+
+
 def _to_text_distributed(  # pylint: disable=unused-argument
     df: pd.DataFrame,
     file_format: str,
@@ -63,16 +109,24 @@ def _to_text_distributed(  # pylint: disable=unused-argument
             path,
         )
 
-    def _datasource_for_format(file_format: str) -> PandasTextDatasource:
-        if file_format == "csv":
-            return PandasCSVDataSource()
-
-        if file_format == "json":
-            return PandasJSONDatasource()
-
-        raise RuntimeError(f"Unknown file format: {file_format}")
+    # Figure out which data source to use, and convert PyArrow parameters if needed
+    try:
+        write_options = _parse_configuration(
+            file_format,
+            s3_additional_kwargs,
+            pandas_kwargs,
+        )
+        can_use_arrow = True
+    except exceptions.InvalidArgument as e:
+        _logger.warning(
+            "PyArrow method unavailable, defaulting to Pandas I/O functions: %s. "
+            "This will result in slower performance of the write operations.",
+            e,
+        )
+        write_options = None
+        can_use_arrow = False
 
-    datasource = _datasource_for_format(file_format)
+    datasource = _datasource_for_format(file_format, can_use_arrow)
 
     mode, encoding, newline = _get_write_details(path=file_path, pandas_kwargs=pandas_kwargs)
     ds.write_datasource(
@@ -87,10 +141,11 @@ def _datasource_for_format(file_format: str) -> PandasTextDatasource:
         dataset_uuid=filename_prefix,
         boto3_session=None,
         s3_additional_kwargs=s3_additional_kwargs,
-        mode=mode,
+        mode="wb" if can_use_arrow else mode,
         encoding=encoding,
         newline=newline,
         pandas_kwargs=pandas_kwargs,
+        write_options=write_options,
     )
 
     return datasource.get_write_paths()
diff --git a/tests/load/test_s3.py b/tests/load/test_s3.py
@@ -127,7 +127,7 @@ def test_s3_delete_objects(path, path2, benchmark_time):
     assert len(wr.s3.list_objects(f"{path2}delete-test*")) == 0
 
 
-@pytest.mark.parametrize("benchmark_time", [240])
+@pytest.mark.parametrize("benchmark_time", [30])
 def test_s3_read_csv_simple(benchmark_time):
     path = "s3://nyc-tlc/csv_backup/yellow_tripdata_2021-0*.csv"
     with ExecutionTimer("elapsed time of wr.s3.read_csv() simple") as timer:
