AWS SDK for pandas' read_parquet using too much memory even when fetching in chunks

Question

We are exploring ways to be able to read large parquet files efficiently in memory (in streams). One of the approaches that seemed promising is [AWS SDK for pandas' read_parquet function](https://aws-sdk-pandas.readthedocs.io/en/stable/stubs/awswrangler.s3.read_parquet.html#awswrangler.s3.read_parquet). The description specifies that using 'chunked=True' or 'chunked=INTEGER' would return smaller chunks as iterable to make reading more memory efficient. 
 
However, when tested with  parquet file of 500mb, the function either times out or goes out of memory. On testing in local, I found out that when we try to receive just one row as a chunk with only one column, the memory usage still peaked at 7GB. So, it seems that behind the scenes this function is still loading a large portion or whole of the file in memory and therefore, it is not doing what its documentation says. 
 
Following is the formatted code we have. Let us know if there is anything wrong with the code or we should make some changes to make it memory efficient. Let us know if you need any other details. Thank you!

```
def get_table_chunks_from_s3_file(app_configs: dict, sqs_values_dict: dict):

bucket = sqs_values_dict["BucketName"]
    key = sqs_values_dict["ObjectKey"]
    boto3_session = app_configs["boto3_session"]

file_path = "s3://" + bucket + "/" + key

dataframes = wr.s3.read_parquet(
        path=file_path, chunked=CHUNK_SIZE_FOR_STREAMING, boto3_session=boto3_session
    )

for dataframe in dataframes:
        yield pyarrow.table(dataframe)
```

re:Post Agent · Answer

The issue you're experiencing with AWS SDK for pandas' read_parquet function using excessive memory even when fetching in chunks is a known limitation. Despite the documentation suggesting that using 'chunked=True' or 'chunked=INTEGER' should return smaller chunks as an iterable for more memory-efficient reading, the function may still load a large portion or the entire file into memory.

To address this problem and make your code more memory efficient, consider the following approaches:

1. Use AWS Glue's custom Parquet writer: AWS Glue provides a custom Parquet writer with performance optimizations for DynamicFrame workflows. This writer computes and modifies the schema dynamically, which can increase performance and potentially reduce memory usage.

2. Adjust Parquet format options: When working with Parquet files, you can fine-tune various options to optimize performance and memory usage. Consider adjusting the following parameters:

- blockSize: Specifies the size in bytes of a row group buffered in memory. The default is 128 MB, but you can reduce this to a smaller value that suits your needs.
- pageSize: Specifies the size in bytes of a page, which is the smallest unit that must be read fully to access a single record. The default is 1 MB, but you can experiment with smaller values.
- compression: Choose an appropriate compression codec (e.g., "snappy", "gzip", or "lzo") that balances between file size and processing speed.

3. Implement manual chunking: Instead of relying on the 'chunked' parameter of read_parquet, you can implement manual chunking by reading specific row groups or columns at a time. This approach gives you more control over memory usage.

4. Use lower-level libraries: Consider using lower-level libraries like pyarrow or fastparquet directly, which offer more fine-grained control over reading Parquet files and memory management.

5. Increase Lambda function memory: If you're running this code in an AWS Lambda function, consider increasing the allocated memory. This will not only provide more memory for processing but also indirectly allocate more CPU capacity, which can improve overall performance.

6. Optimize your code: Ensure that you're not holding onto unnecessary data in memory. Process and release data as soon as possible within your function.

By implementing these strategies, you should be able to reduce memory usage and improve the efficiency of reading large Parquet files from S3. Remember to test thoroughly with various file sizes and configurations to find the optimal setup for your specific use case.  
**Sources**  
[Troubleshooting Lambda configurations - AWS Lambda](https://docs.aws.amazon.com/lambda/latest/operatorguide/configurations.html)  
[Using the Parquet format in AWS Glue - AWS Glue](https://docs.aws.amazon.com/glue/latest/dg/aws-glue-programming-etl-format-parquet-home.html)

AWS SDK for pandas' read_parquet using too much memory even when fetching in chunks

Add your answer

Relevant content