Aws-data-wrangler: Write multiple parquet files from a single dataframe defined by max file size

[chrispruitt]

The Problem
I was unsuccessful in finding a way to write a single dataframe to multiple parquet files using the s3.to_parquet() method. Currently, it seems to write one parquet file which could slow down Athena queries.

Possible Solution
It would be great to have an option for the "max parquet file size" when using s3.to_parquet(). So instead of creating one large parquet file, we could create many smaller parquet files that would help optimize Athena queries.

My reasoning
I believe Athena, behind the scenes, uses the number of files to split the query workload across different nodes.

[rparthas]

Can I work on this particular request ?

[igorborgest]

@rparthas yep, it would be very welcome! Just make sure to checkout from our dev branch and then open a pull request against the same.

[igorborgest]

@chrispruitt @rparthas

We already have this feature done to release. Could you both test it directly from our dev branch before the official release? The ideia is publish the feature in the version 1.8.0 next weekend.

pip install git+https://github.com/awslabs/aws-data-wrangler.git@dev

e.g:

wr.s3.to_parquet(
    df=df,
    path="s3://.../myfile.parquet",
    max_rows_by_file=1_000_000,
)

P.S. Thanks a lot for your great work @rparthas, I've basically changed the feature to use max_rows_by_file instead of the file size itself because it is pretty tough to try to guess what will be the size of parquet file serialised on s3 and I don't think that the PyArrow table size will be a good guide for user.

P.P.S. In addition, is expected that writes will be faster when chunked because we are parallelising all files uploads.

[igorborgest]

Hi @chrispruitt @rparthas !

It is totally off-topic, but we are stating a _"Who uses AWS Data Wrangler?"_ section. So feel free to add yourself if you want 😄 .

[igorborgest]

Released on version 1.8.0.