apache / apache/arrow

pyarrow may use a lot of memory to load a dataframe from parquet

Open
#32,670 10 comments 0 reactions 0 assignees View on GitHub
Component: Parquet Component: Python Type: bug
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

When a pandas dataframe is loaded from a parquet file using `{}pyarrow.parquet.read_table{`}, the memory usage may grow a lot more than what should be needed to load the dataframe, and it's not freed until the dataframe is deleted.

The problem is evident when the dataframe has a {**}column containing lists or numpy arrays{**}, while it seems absent (or not noticeable) if the column contains only integer or floats.

I'm attaching a simple script to reproduce the issue, and a graph created with memory-profiler showing the memory usage.

In this example, the dataframe created with pandas needs around 1.2 GB, but the memory usage after loading it from parquet is around 16 GB.

The items of the column are created as numpy arrays and not lists, to be consistent with the types loaded from parquet (pyarrow produces numpy arrays and not lists).

 
```python

import gc
import time
import numpy as np
import pandas as pd
import pyarrow
import pyarrow.parquet
import psutil

def pyarrow_dump(filename, df, compression="snappy"):
table = pyarrow.Table.from_pandas(df)
pyarrow.parquet.write_table(table, filename, compression=compression)

def pyarrow_load(filename):
table = pyarrow.parquet.read_table(filename)
return table.to_pandas()

def print_mem(msg, start_time=time.monotonic(), process=psutil.Process()):
# gc.collect()
current_time = time.monotonic() - start_time
rss = process.memory_info().rss / 2 ** 20
print(f"{msg:>3} time:{current_time:>10.1f} rss:{rss:>10.1f}")

if __name__ == "__main__":
print_mem(0)

rows = 5000000
df = pd.DataFrame({"a": [np.arange(10) for i in range(rows)]})
print_mem(1)

pyarrow_dump("example.parquet", df)
print_mem(2)

del df
print_mem(3)
time.sleep(3)
print_mem(4)

df = pyarrow_load("example.parquet")
print_mem(5)
time.sleep(3)
print_mem(6)

del df
print_mem(7)
time.sleep(3)
print_mem(8)
```
Run with memory-profiler:
```bash

mprof run --multiprocess python test_pyarrow.py
```
Output:
```java

mprof: Sampling memory every 0.1s
running new process
0 time: 0.0 rss: 135.4
1 time: 4.9 rss: 1252.2
2 time: 7.1 rss: 1265.0
3 time: 7.5 rss: 760.2
4 time: 10.7 rss: 758.9
5 time: 19.6 rss: 16745.4
6 time: 22.6 rss: 16335.4
7 time: 22.9 rss: 15833.0
8 time: 25.9 rss: 955.0
```

**Environment**: linux
**Reporter**: [Gianluca Ficarelli](https://issues.apache.org/jira/browse/ARROW-17399)
**Watchers**: [Rok Mihevc](https://issues.apache.org/jira/browse/ARROW-17399) / @rok
#### Original Issue Attachments:
- [memory-profiler.png](https://issues.apache.org/jira/secure/attachment/13048063/memory-profiler.png)

**Note**: *This issue was originally created as [ARROW-17399](https://issues.apache.org/jira/browse/ARROW-17399). Please see the [migration documentation](https://github.com/apache/arrow/issues/14542) for further details.*

Contributor guide

Open the contributing guide

Research direction

Run the attached test_pyarrow.py reproducer, focusing on pyarrow.parquet.read_table and table.to_pandas() in pyarrow_load. Compare memory usage for columns containing NumPy arrays with the reported 1.2 GB versus 16 GB results, and confirm that loading and deleting the dataframe no longer retains the excess memory.

Written by the indexing model from the issue text.

Assessment

Tech stack
pandas, python
Domain
data
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.