[Python][Parquet] Memory Leak when writing
- Dominant language
- C++
- Stars
- 17.1k
- Forks
- 4.3k
- Avg merge
- 3d 13h
- Merged PRs (30d)
- 88
Description
## Summary
After investigating an uwsgi worker crash we found that a memory leak of our parquet export was the cause.
We use python 3.9 with pyarrow 17.0.0 but this is reproductible with the following code and also with python 3.12.
We found that the problem comes from `writer.writer.write_table(table, row_group_size=1)` (row_group_size=1 for the example).
If you have any advice on how to reduce or fix this memory leak it will be very appreciated 🙏
## How to reproduce?
Install: `pip install pyarrow memory_profiler`
Then use the following script to generate 100 000 records with an empty `{"a": ""}`. You will show that the RSS is increasing up to 90mo for only 100 000 records before closing the writer and still stay at +75mo after closing it!
```python
import io
import pyarrow as pa
import pyarrow.parquet as pq
import psutil
from memory_profiler import profile
# change it to mimalloc on Mac (this can run in a debian container)
pool = pa.jemalloc_memory_pool()
pa.jemalloc_set_decay_ms(0)
pa.set_memory_pool(pool)
previous_rss = 0
first_rss = 0
@profile
def test(previous_rss, first_rss):
def yield_batch(count: int, schema):
for i in range(count):
yield i, None # pa.RecordBatch.from_pylist([{"a": "1"}], schema=schema)
buf = io.BytesIO()
buffer = pa.output_stream(buf, compression=None, buffer_size=None)
schema = pa.schema([pa.field("a", pa.string())])
writer = pq.ParquetWriter(
buffer,
schema=schema,
compression='none',
write_statistics=False,
write_page_index=False,
write_page_checksum=False,
use_dictionary=False,
store_schema=True,
)
bytes_length = 0
table = pa.Table.from_pylist([{"a": ""}], schema=schema)
for i, batch in yield_batch(100000, schema):
writer.writer.write_table(table, row_group_size=1)
buf.seek(0)
val = buf.read()
bytes_length += len(val)
del val
buf.flush()
buf.truncate(0)
buf.seek(0)
if i % 1_000 == 0:
p = psutil.Process().memory_info()
if previous_rss == 0:
previous_rss = p.rss
first_rss = p.rss
diff = p.rss - previous_rss
print((
f"Chunk: {i}; RSS: {p.rss / 1_000_000}; "
f"(mb) Pool size:{pool.bytes_allocated()}; "
f"Diff: {diff}; +{(p.rss - first_rss) / 1_000_000} (mb);"
))
previous_rss = p.rss
writer.close()
buf.seek(0)
val = buf.read()
bytes_length += len(val)
del val
buf.flush()
buf.truncate(0)
buf.seek(0)
print(f"Export size: {bytes_length / 1_000_000} mb")
return previous_rss, first_rss
if __name__=='__main__':
previous_rss, first_rss = test(previous_rss, first_rss)
p = psutil.Process().memory_info()
if previous_rss == 0:
previous_rss = p.rss
first_rss = p.rss
diff = p.rss - previous_rss
print((
f"RSS: {p.rss / 1_000_000}; "
f"(mb) Pool size:{pool.bytes_allocated()}; "
f"Diff: {diff}; +{(p.rss - first_rss) / 1_000_000} (mb);"
))
previous_rss = p.rss
```
Output:
```bash
Chunk: 0; RSS: 90.619904; (mb) Pool size:384; Diff: 0; +0.0 (mb);
Chunk: 1000; RSS: 91.25888; (mb) Pool size:384; Diff: 638976; +0.638976 (mb);
Chunk: 2000; RSS: 91.930624; (mb) Pool size:384; Diff: 671744; +1.31072 (mb);
Chunk: 3000; RSS: 93.650944; (mb) Pool size:384; Diff: 1720320; +3.03104 (mb);
Chunk: 4000; RSS: 94.388224; (mb) Pool size:384; Diff: 737280; +3.76832 (mb);
Chunk: 5000; RSS: 95.748096; (mb) Pool size:384; Diff: 1359872; +5.128192 (mb);
Chunk: 6000; RSS: 96.534528; (mb) Pool size:384; Diff: 786432; +5.914624 (mb);
Chunk: 7000; RSS: 97.271808; (mb) Pool size:384; Diff: 737280; +6.651904 (mb);
Chunk: 8000; RSS: 98.05824; (mb) Pool size:384; Diff: 786432; +7.438336 (mb);
Chunk: 9000; RSS: 99.61472; (mb) Pool size:384; Diff: 1556480; +8.994816 (mb);
Chunk: 10000; RSS: 100.450304; (mb) Pool size:384; Diff: 835584; +9.8304 (mb);
Chunk: 11000; RSS: 101.220352; (mb) Pool size:384; Diff: 770048; +10.600448 (mb);
Chunk: 12000; RSS: 102.023168; (mb) Pool size:384; Diff: 802816; +11.403264 (mb);
Chunk: 13000; RSS: 102.760448; (mb) Pool size:384; Diff: 737280; +12.140544 (mb);
Chunk: 14000; RSS: 103.514112; (mb) Pool size:384; Diff: 753664; +12.894208 (mb);
Chunk: 15000; RSS: 104.316928; (mb) Pool size:384; Diff: 802816; +13.697024 (mb);
Chunk: 16000; RSS: 105.070592; (mb) Pool size:384; Diff: 753664; +14.450688 (mb);
Chunk: 17000; RSS: 107.413504; (mb) Pool size:384; Diff: 2342912; +16.7936 (mb);
Chunk: 18000; RSS: 108.232704; (mb) Pool size:384; Diff: 819200; +17.6128 (mb);
Chunk: 19000; RSS: 109.002752; (mb) Pool size:384; Diff: 770048; +18.382848 (mb);
Chunk: 20000; RSS: 109.7728; (mb) Pool size:384; Diff: 770048; +19.152896 (mb);
Chunk: 21000; RSS: 110.51008; (mb) Pool size:384; Diff: 737280; +19.890176 (mb);
Chunk: 22000; RSS: 111.312896; (mb) Pool size:384; Diff: 802816; +20.692992 (mb);
Chunk: 23000; RSS: 112.082944; (mb) Pool size:384; Diff: 770048; +21.46304 (mb);
Chunk: 24000; RSS: 112.88576; (mb) Pool size:384; Diff: 802816; +22.265856 (mb);
Chunk: 25000; RSS: 113.639424; (mb) Pool size:384; Diff: 753664; +23.01952 (mb);
Chunk: 26000; RSS: 114.393088; (mb) Pool size:384; Diff: 753664; +23.773184 (mb);
Chunk: 27000; RSS: 115.146752; (mb) Pool size:384; Diff: 753664; +24.526848 (mb);
Chunk: 28000; RSS: 115.982336; (mb) Pool size:384; Diff: 835584; +25.362432 (mb);
Chunk: 29000; RSS: 116.752384; (mb) Pool size:384; Diff: 770048; +26.13248 (mb);
Chunk: 30000; RSS: 117.506048; (mb) Pool size:384; Diff: 753664; +26.886144 (mb);
Chunk: 31000; RSS: 118.276096; (mb) Pool size:384; Diff: 770048; +27.656192 (mb);
Chunk: 32000; RSS: 119.062528; (mb) Pool size:384; Diff: 786432; +28.442624 (mb);
Chunk: 33000; RSS: 122.994688; (mb) Pool size:384; Diff: 3932160; +32.374784 (mb);
Chunk: 34000; RSS: 123.748352; (mb) Pool size:384; Diff: 753664; +33.128448 (mb);
Chunk: 35000; RSS: 124.551168; (mb) Pool size:384; Diff: 802816; +33.931264 (mb);
Chunk: 36000; RSS: 125.304832; (mb) Pool size:384; Diff: 753664; +34.684928 (mb);
Chunk: 37000; RSS: 126.042112; (mb) Pool size:384; Diff: 737280; +35.422208 (mb);
Chunk: 38000; RSS: 126.861312; (mb) Pool size:384; Diff: 819200; +36.241408 (mb);
Chunk: 39000; RSS: 127.647744; (mb) Pool size:384; Diff: 786432; +37.02784 (mb);
Chunk: 40000; RSS: 128.434176; (mb) Pool size:384; Diff: 786432; +37.814272 (mb);
Chunk: 41000; RSS: 129.253376; (mb) Pool size:384; Diff: 819200; +38.633472 (mb);
Chunk: 42000; RSS: 129.990656; (mb) Pool size:384; Diff: 737280; +39.370752 (mb);
Chunk: 43000; RSS: 130.793472; (mb) Pool size:384; Diff: 802816; +40.173568 (mb);
Chunk: 44000; RSS: 131.547136; (mb) Pool size:384; Diff: 753664; +40.927232 (mb);
Chunk: 45000; RSS: 132.317184; (mb) Pool size:384; Diff: 770048; +41.69728 (mb);
Chunk: 46000; RSS: 133.103616; (mb) Pool size:384; Diff: 786432; +42.483712 (mb);
Chunk: 47000; RSS: 133.873664; (mb) Pool size:384; Diff: 770048; +43.25376 (mb);
Chunk: 48000; RSS: 134.643712; (mb) Pool size:384; Diff: 770048; +44.023808 (mb);
Chunk: 49000; RSS: 135.380992; (mb) Pool size:384; Diff: 737280; +44.761088 (mb);
Chunk: 50000; RSS: 136.15104; (mb) Pool size:384; Diff: 770048; +45.531136 (mb);
Chunk: 51000; RSS: 136.97024; (mb) Pool size:384; Diff: 819200; +46.350336 (mb);
Chunk: 52000; RSS: 137.773056; (mb) Pool size:384; Diff: 802816; +47.153152 (mb);
Chunk: 53000; RSS: 138.510336; (mb) Pool size:384; Diff: 737280; +47.890432 (mb);
Chunk: 54000; RSS: 139.296768; (mb) Pool size:384; Diff: 786432; +48.676864 (mb);
Chunk: 55000; RSS: 140.066816; (mb) Pool size:384; Diff: 770048; +49.446912 (mb);
Chunk: 56000; RSS: 140.853248; (mb) Pool size:384; Diff: 786432; +50.233344 (mb);
Chunk: 57000; RSS: 141.606912; (mb) Pool size:384; Diff: 753664; +50.987008 (mb);
Chunk: 58000; RSS: 142.409728; (mb) Pool size:384; Diff: 802816; +51.789824 (mb);
Chunk: 59000; RSS: 143.163392; (mb) Pool size:384; Diff: 753664; +52.543488 (mb);
Chunk: 60000; RSS: 143.93344; (mb) Pool size:384; Diff: 770048; +53.313536 (mb);
Chunk: 61000; RSS: 144.703488; (mb) Pool size:384; Diff: 770048; +54.083584 (mb);
Chunk: 62000; RSS: 145.522688; (mb) Pool size:384; Diff: 819200; +54.902784 (mb);
Chunk: 63000; RSS: 146.276352; (mb) Pool size:384; Diff: 753664; +55.656448 (mb);
Chunk: 64000; RSS: 147.030016; (mb) Pool size:384; Diff: 753664; +56.410112 (mb);
Chunk: 65000; RSS: 147.78368; (mb) Pool size:384; Diff: 753664; +57.163776 (mb);
Chunk: 66000; RSS: 154.91072; (mb) Pool size:384; Diff: 7127040; +64.290816 (mb);
Chunk: 67000; RSS: 155.680768; (mb) Pool size:384; Diff: 770048; +65.060864 (mb);
Chunk: 68000; RSS: 156.483584; (mb) Pool size:384; Diff: 802816; +65.86368 (mb);
Chunk: 69000; RSS: 157.2864; (mb) Pool size:384; Diff: 802816; +66.666496 (mb);
Chunk: 70000; RSS: 158.072832; (mb) Pool size:384; Diff: 786432; +67.452928 (mb);
Chunk: 71000; RSS: 158.826496; (mb) Pool size:384; Diff: 753664; +68.206592 (mb);
Chunk: 72000; RSS: 159.629312; (mb) Pool size:384; Diff: 802816; +69.009408 (mb);
Chunk: 73000; RSS: 160.39936; (mb) Pool size:384; Diff: 770048; +69.779456 (mb);
Chunk: 74000; RSS: 161.13664; (mb) Pool size:384; Diff: 737280; +70.516736 (mb);
Chunk: 75000; RSS: 161.988608; (mb) Pool size:384; Diff: 851968; +71.368704 (mb);
Chunk: 76000; RSS: 162.725888; (mb) Pool size:384; Diff: 737280; +72.105984 (mb);
Chunk: 77000; RSS: 163.495936; (mb) Pool size:384; Diff: 770048; +72.876032 (mb);
Chunk: 78000; RSS: 164.265984; (mb) Pool size:384; Diff: 770048; +73.64608 (mb);
Chunk: 79000; RSS: 165.0688; (mb) Pool size:384; Diff: 802816; +74.448896 (mb);
Chunk: 80000; RSS: 165.855232; (mb) Pool size:384; Diff: 786432; +75.235328 (mb);
Chunk: 81000; RSS: 166.641664; (mb) Pool size:384; Diff: 786432; +76.02176 (mb);
Chunk: 82000; RSS: 167.411712; (mb) Pool size:384; Diff: 770048; +76.791808 (mb);
Chunk: 83000; RSS: 168.214528; (mb) Pool size:384; Diff: 802816; +77.594624 (mb);
Chunk: 84000; RSS: 168.935424; (mb) Pool size:384; Diff: 720896; +78.31552 (mb);
Chunk: 85000; RSS: 169.689088; (mb) Pool size:384; Diff: 753664; +79.069184 (mb);
Chunk: 86000; RSS: 170.459136; (mb) Pool size:384; Diff: 770048; +79.839232 (mb);
Chunk: 87000; RSS: 171.311104; (mb) Pool size:384; Diff: 851968; +80.6912 (mb);
Chunk: 88000; RSS: 172.081152; (mb) Pool size:384; Diff: 770048; +81.461248 (mb);
Chunk: 89000; RSS: 172.867584; (mb) Pool size:384; Diff: 786432; +82.24768 (mb);
Chunk: 90000; RSS: 173.654016; (mb) Pool size:384; Diff: 786432; +83.034112 (mb);
Chunk: 91000; RSS: 174.40768; (mb) Pool size:384; Diff: 753664; +83.787776 (mb);
Chunk: 92000; RSS: 175.177728; (mb) Pool size:384; Diff: 770048; +84.557824 (mb);
Chunk: 93000; RSS: 175.931392; (mb) Pool size:384; Diff: 753664; +85.311488 (mb);
Chunk: 94000; RSS: 176.70144; (mb) Pool size:384; Diff: 770048; +86.081536 (mb);
Chunk: 95000; RSS: 177.487872; (mb) Pool size:384; Diff: 786432; +86.867968 (mb);
Chunk: 96000; RSS: 178.307072; (mb) Pool size:384; Diff: 819200; +87.687168 (mb);
Chunk: 97000; RSS: 179.109888; (mb) Pool size:384; Diff: 802816; +88.489984 (mb);
Chunk: 98000; RSS: 179.879936; (mb) Pool size:384; Diff: 770048; +89.260032 (mb);
Chunk: 99000; RSS: 180.649984; (mb) Pool size:384; Diff: 770048; +90.03008 (mb);
Export size: 12.596596 mb
Filename: /Users/xxxxxxx/Documents/github/xxxxxx/xxxxxxx.py
Line # Mem usage Increment Occurrences Line Contents
=============================================================
15 42.9 MiB 42.9 MiB 1 @profile
16 def test(previous_rss, first_rss):
17
18 85.9 MiB 0.0 MiB 2 def yield_batch(count: int, schema):
19 173.0 MiB -0.2 MiB 100001 for i in range(count):
20 173.0 MiB -0.3 MiB 200000 yield i, None # pa.RecordBatch.from_pylist([{"a": "1"}], schema=schema)
21
22 42.9 MiB 0.0 MiB 1 buf = io.BytesIO()
23 42.9 MiB 0.1 MiB 1 buffer = pa.output_stream(buf, compression=None, buffer_size=None)
24 43.0 MiB 0.1 MiB 1 schema = pa.schema([pa.field("a", pa.string())])
25 44.0 MiB 1.0 MiB 2 writer = pq.ParquetWriter(
26 43.0 MiB 0.0 MiB 1 buffer,
27 43.0 MiB 0.0 MiB 1 schema=schema,
28 43.0 MiB 0.0 MiB 1 compression='none',
29 43.0 MiB 0.0 MiB 1 write_statistics=False,
30 43.0 MiB 0.0 MiB 1 write_page_index=False,
31 43.0 MiB 0.0 MiB 1 write_page_checksum=False,
32 43.0 MiB 0.0 MiB 1 use_dictionary=False,
33 43.0 MiB 0.0 MiB 1 store_schema=True,
34 )
35 44.0 MiB 0.0 MiB 1 bytes_length = 0
36 85.9 MiB 41.9 MiB 1 table = pa.Table.from_pylist([{"a": ""}], schema=schema)
37
38 173.0 MiB -0.2 MiB 100001 for i, batch in yield_batch(100000, schema):
39 173.0 MiB 86.9 MiB 100000 writer.writer.write_table(table, row_group_size=1)
40
41 173.0 MiB -0.8 MiB 100000 buf.seek(0)
42 173.0 MiB -0.2 MiB 100000 val = buf.read()
43 173.0 MiB -0.2 MiB 100000 bytes_length += len(val)
44 173.0 MiB -0.2 MiB 100000 del val
45 173.0 MiB -0.2 MiB 100000 buf.flush()
46 173.0 MiB -0.2 MiB 100000 buf.truncate(0)
47 173.0 MiB -0.1 MiB 100000 buf.seek(0)
48
49 173.0 MiB -0.2 MiB 100000 if i % 1_000 == 0:
50 172.3 MiB 0.0 MiB 100 p = psutil.Process().memory_info()
51 172.3 MiB 0.0 MiB 100 if previous_rss == 0:
52 86.4 MiB 0.0 MiB 1 previous_rss = p.rss
53 86.4 MiB 0.0 MiB 1 first_rss = p.rss
54 172.3 MiB 0.0 MiB 100 diff = p.rss - previous_rss
55 172.3 MiB 0.0 MiB 200 print((
56 172.3 MiB 0.0 MiB 400 f"Chunk: {i}; RSS: {p.rss / 1_000_000}; "
57 172.3 MiB 0.0 MiB 100 f"(mb) Pool size:{pool.bytes_allocated()}; "
58 172.3 MiB 0.0 MiB 200 f"Diff: {diff}; +{(p.rss - first_rss) / 1_000_000} (mb);"
59 ))
60 172.3 MiB 0.0 MiB 100 previous_rss = p.rss
61
62 226.0 MiB 53.0 MiB 1 writer.close()
63
64 226.0 MiB 0.0 MiB 1 buf.seek(0)
65 226.0 MiB 0.0 MiB 1 val = buf.read()
66 226.0 MiB 0.0 MiB 1 bytes_length += len(val)
67 226.0 MiB 0.0 MiB 1 del val
68 226.0 MiB 0.0 MiB 1 buf.flush()
69 226.0 MiB 0.0 MiB 1 buf.truncate(0)
70 226.0 MiB 0.0 MiB 1 buf.seek(0)
71
72 226.0 MiB 0.0 MiB 1 print(f"Export size: {bytes_length / 1_000_000} mb")
73 226.0 MiB 0.0 MiB 1 return previous_rss, first_rss
RSS: 164.806656; (mb) Pool size:0; Diff: -15843328; +74.186752 (mb);
```
### Component(s)
Parquet, Python
Contributor guide
Assessment
This issue has not been assessed yet.