Slow performance when using read_parquet from s3

[konradsemsch]

Hi,

I would like to open an issue as we have seen quite unsatisfying performance using the read_parquet function. This is our setup and data below:

• data is in S3, there are 1164 individual date-time prefixes under the main folder, and the total size of all files is barely 25.6 MB. So quite a lot of small individual files organized by individual prefixes by date-time
• they way we gather these files is by passing path: s3://.../ and using the partition_filter. The function call looks like this:

wr.s3.read_parquet(
                path,
                dataset=True,
                partition_filter=filter(),
            )

I've run a couple of tests to verify whether there would be any speed improvement if I passed a list of prefixes for the function to combine instead of using the partition_filter but the gain was marginal. Enabling use_threads=True gave no improvement. Overall it takes around 13 minutes to collect all files... this is just too long. Downloading them with aws sync takes a few seconds.

Our main use case for operating on streams is in AWS Batch. We have some data loaders that use the data wrangler when we train our ML model in AWS Batch. We realized after some time that the main contributor to an extended training time, is the part where the data is collected from AWS using the data wrangler (primarily the wr.s3.read_parquet). Please also note that we're not taking of big data here. Most of our use cases is like described above.

At the moment we're wondering whether this can be optimized or if we should move away from the streaming approach, and simply download the data on the container for model training. Could you give some advice? What's your take on that?

[maxispeicher]

After having a first look I think one of the major slowdowns is that even with use_threads=True each file is downloaded with multiple threads but the different files are still loaded sequentially.
Some performance numbers from a first test:

import time
import pandas as pd
import awswrangler as wr

i = 0
data = []
for day in range(1, 10):
    for hour in range(24):
        for _ in range(1000):
            data.append({"value": i, "year": 2021, "month": 3, "day": day, "hour": hour})
            i = i+1
df = pd.DataFrame(data)

s3_location = "s3://my_s3_location"
wr.s3.to_parquet(df, path=my_s3_location, partition_cols=["year", "month", "day", "hour"])

start_time = time.perf_counter()
df2 = wr.s3.read_parquet(s3_location)
print(f"Loading complete data took {time.perf_counter() - start_time} seconds.")
# >>> Loading complete data took 260 seconds.

start_time = time.perf_counter()
df3 = wr.s3.read_parquet(s3_location, partition_filter=lambda x: x["day"] == 1, dataset=True)
print(f"Loading one day took {time.perf_counter() - start_time} seconds.")
# >>> Loading one day took 29.34 seconds.

Then I tried to do a first draft implementation using concurrent.futures.ThreadPoolExecutor for downloading files in parallel instead of https://github.com/awslabs/aws-data-wrangler/blob/main/awswrangler/s3/_read_parquet.py#L633:

with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
    futures = []
    for p in paths:
        partial_func = functools.partial(_read_parquet, path=p)
        futures.append(executor.submit(partial_func, **args))
results = []
for completed in concurrent.futures.as_completed(futures):
    results.append(completed.result())
return _union(dfs=results, ignore_index=ignore_index)

resulting in

start_time = time.perf_counter()
df3 = wr.s3.read_parquet(s3_location, partition_filter=lambda x: x["day"] == 1, dataset=True)
print(f"Loading one day with threads took {time.perf_counter() - start_time} seconds.")
# >>> Loading one day with threads took 8.7 seconds.

[gvermillion]

Just dropping some support behind having awswrangler load files in parallel. We have a similar use case to that described above. We use AWS Batch to run large scale simulations, and awswrangler is used to pull the necessary input data sets down for cleaning, prepping, etc. Initially, the large input data sets are just buckets of CSV files (no partitioning logic/structure), and that initial load is a huge bottleneck.

[maxispeicher]

I've just pushed a very basic draft how the parallel downloading could be enabled. However just to be on the save side if it only increases performance for my synthetic example, it would be nice if you could test it for your scenarios (@konradsemsch, @gvermillion). Right it uses the number of available cores for downloading the files in parallel, but if the solution seems to work out I think we can also directly exhibit this parameter aswell.
If you want to test it, you can install it with

pip uninstall awswrangler -y
pip install git+https://github.com/maxispeicher/aws-data-wrangler.git@enable_parallel_s3_downloads

Right now it should support parallel downloads for read_parquet, read_csv, read_json, read_fwf (however only read_parquet was tested by me until now).

[konradsemsch]

Hi @maxispeicher! Thanks for quick feedback and proposing a solution already :) However, we probably would need to wait anyways until this change is merged, as we need this for our production workloads...

[gvermillion]

@maxspeicher, I've slotted time to test this this week. I'll report back.

[konradsemsch]

@maxispeicher - I've tried out our implementation and I see a massive performance boost. Processing time of reading the files got down in our case by a factor of 4: from 8.5 mins to 2.2 mins! Good job! 😉 🍺

[maxispeicher]

@konradsemsch Happy to hear that 🙂
I've now also created a PR adressing the issue. It is now also possible to directly set the number of used threads via the use_threads parameter which might also be able to increase your speed performance even further 😉.
e.g.

start_time = time.perf_counter()
df3 = wr.s3.read_parquet(s3_location, partition_filter=lambda x: x["day"] == 1, dataset=True, use_threads=16)
print(f"Loading one day with threads took {time.perf_counter() - start_time} seconds.")
# >>> Loading one day with threads took 4.47 seconds. (instead of ~9 seconds with just use_threads=True)