(@scale): Reading a large number of small S3 objects is slow and might eventually fail

[jaidisido]

Describe the bug

In distributed mode, reading a large number of small S3 objects (e.g 20M files) is slow and might eventually fail.

This is caused by the list objects call which is not currently parallelised and represents a bottleneck.

As a side note, no information is surfaced to the user during the list object call, meaning they don't have visibility into why the job is hanging. Additional logging would improve that.

How to Reproduce

Setup

• 20 million files with .csv extension
• each file contains 5 lines with 4 columns

Script

import awswrangler as wr
logging.getLogger("awswrangler").setLevel(logging.DEBUG)

df = wr.s3.read_csv("s3://test/small-files/input/20m-csv-partition/")
print(df.head())

Logs

2023-02-02 01:00:33,152 - awswrangler._config - DEBUG - Applying default config argument verify with value None.
2023-02-02 01:00:33,155 - botocore.credentials - INFO - Found credentials in shared credentials file: ~/.aws/credentials
2023-02-02 01:00:33,268 - awswrangler.s3._list - DEBUG - args: {'Bucket': 'test', 'Prefix': 'small-files/input/20m-csv-partition/', 'PaginationConfig': {'PageSize': 1000}}
2023-02-02 01:00:33,550 - awswrangler.s3._list - DEBUG - Skipping empty file: s3://test/small-files/input/20m-csv-partition/
2023-02-02 02:00:16,026 - __main__ - INFO -   File "<stdin>", line 8, in <module>
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_read_text.py", line 281, in read_csv
    return _read_text_format(
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_read_text.py", line 91, in _read_text_format
    paths: List[str] = _path2list(
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_list.py", line 31, in _path2list
    paths: List[str] = list_objects(  # type: ignore
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_list.py", line 358, in list_objects
    return [path for paths in result_iterator for path in paths]
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_list.py", line 358, in <listcomp>
    return [path for paths in result_iterator for path in paths]
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/awswrangler/s3/_list.py", line 110, in _list_objects
    for page in response_iterator:  # pylint: disable=too-many-nested-blocks
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/botocore/paginate.py", line 269, in __iter__
    response = self._make_request(current_kwargs)
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/botocore/paginate.py", line 357, in _make_request
    return self._method(**current_kwargs)
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/botocore/client.py", line 530, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/opt/amazon/python3.9-ray/lib/python3.9/site-packages/botocore/client.py", line 960, in _make_api_call
    raise error_class(parsed_response, operation_name)
botocore.exceptions.ClientError: An error occurred (ExpiredToken) when calling the ListObjectsV2 operation: The provided token has expired.
2023-02-02 02:00:16,026 - __main__ - WARNING - botocore.exceptions.ClientError: An error occurred (ExpiredToken) when calling the ListObjectsV2 operation: The provided token has expired.

Script eventually fails with an expired token

Expected behavior

No response

Your project

No response

Screenshots

No response

OS

Unix

Python version

3.9

AWS SDK for pandas version

3.0.0rc2

Tasks

• Create 20M CSV files in S3 bucket (APG)
• Test with ray.data.read_csv and compare performance
• Consider delegating path listing to Ray or see if we can replicate the same logic
• Explore parallelising S3 list objects call

Ray implementation: https://github.com/ray-project/ray/blob/master/python/ray/data/datasource/file_meta_provider.py#L189

[jaidisido]

Summary of findings:

1. The Expired Token exception is related to this issue which is a known boto3 limitation. Unless an assume role provider is used, the boto3 session expires when the credentials are refreshed. Credentials rotation on a Glue on Ray job is set to 1h.

Findings from @kukushking:

2. Before executing a read, Ray prepares the list of read tasks, which involves expanding S3 paths and collecting metadata. This task is not parallelised by default in vanilla ray.data and SDK for pandas implementations, meaning it only consumes a single CPU. However, setting the parallelism argument to a large value (e.g. 1,000) helps parallelise these tasks. By default, parallelism is set to -1 (i.e. let Ray infer based on the task). Ray does not set a high value in this case because read tasks are I/O intensive (not CPU-bounded).

3. SDK for pandas performs an S3 list before passing a list of paths to Ray. In the current implementation, Ray expands the paths a second time, resulting in wasteful s3 calls. The issue is mitigated by using FastFileMetadataProvider which skips paths expansion.

4. In addition, a new ArrowCSVDatasource was introduced in SDK for pandas after rc2 which has better read performance than PandasCSVDataSource.

By using the FastFileMetadataProvider and ArrowCSVDatasource, the performance of SDK for pandas is on-par with the vanilla Ray implementation based on a test on 1M small CSV files.

5. Parallelising the S3 list call could be considered. However it would only be beneficial if the structure of the S3 bucket has multiple prefixes. An argument can be exposed to enable the user to request parallelisation of the call. By default it would be set to False.

[jjacobj84]

@jaidisido would this optimization be done only for CSV reads or we are addressing this for all read formats.

[kukushking]

@jjacobj84 this is done for all text-based formats, parquet is a bit special. For parquet, we might consider @jaidisido 's suggestion in #2023 . It looks like bulk parquet reads differ from regular parquet in that 1) skips paths expansion 2) avoids using arrows parquet datasets which are bottleneck albeit they provide unified schema