read_parquet method in chunks is reading entire dataset into memory

[chariottrider]

Using wr.read_parquet(path, chunked=10) vs. wr.read_parquet(path, chunked=10000000) on an EC2 instance is consuming the same amount of memory and is basically the memory needed to hold the the entire dataset in memory.

wr.read_csv(chunksize=10) does not do this. The memory usage as you iterate on chunks is very small as expected and makes it possible to iterate on very large datasets.

pyarrow 3.0 supports iter_batches() method, which works much like wr_read_csv() in that the memory impact is very low and yet iterating in chunks is very fast.

I feel the current implementation of read_parquet with specifying an int chunksize is misleading as it seems like it should act like read_csv with chunksize and it doesn't. Not even close. Perhaps consider using pyarrow for this so that we can iterate on very large datasets efficiently.

Here's an example of working code I used for iterating on one parquet file. This could easily be wrapped to iterate on multiple files in an s3 location.

import pyarrow.parquet as pq
from smart_open import open

path_like_object = open("s3://bucket/key_to_single_file.parquet", 'rb')
#you can use smart_open, s3fs, boto or whatever lib you want to open the path as the ParquetFile class below will accept a file-like object

pfile = pq.ParquetFile(path_like_object)

for table in pfile.iter_batches(batch_size=100):
    df = table.to_pandas()
    print (df)

[maxispeicher]

Hey @chariottrider,

thanks for reporting. I've crated a PR which should use iter_batches when it is available but keeps the old behavior for pyarrow < 3.0.0. If you want you can give it a try:

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

[chariottrider]

@maxispeicher - Now that was fast!!

I tried it out and it's perfect! Great speed and keeps the memory consumption inline while iterating over our large features store. Having this baked into Wrangler will really help our data scientists who are already familiar with Wrangler.

Thanks a ton! Any estimate as to when this might be released with a regular version?

[maxispeicher]

@chariottrider Happy to hear that it works for you 🙂
I just pushed another iteration, which hopefully better integrates with the existing code. But from a functional perspective it should have stayed almost the same. Only Difference: If you e.g. had 5 parquet files with 10 lines and read with a chunk size of 50, the first version would return a generator with 5 dfs of length 10, whereas the newly pushed version only has one df with length 50.

Regarding a release there's nothing I can say for sure, but in general a new release will be available roughly each month. But probably @jaidisido can you tell more about it 🙂.

[chariottrider]

Many thanks @maxispeicher

[jaidisido]

Thank you @chariottrider for raising it and @maxispeicher for the fix. Releases are usually monthly indeed. 2.7.0 was release on April 15th

[jaidisido]

Now covered in 2.8.0.