Memory Leak when Importing Parquet File with PyArrow Engine in Pandas

[Voltagabbana]

Description

We've identified a memory leak when importing Parquet files into Pandas DataFrames using the PyArrow engine. The issue occurs specifically during the conversion from Arrow to Pandas objects, as memory is not released even after deleting the DataFrame and invoking garbage collection.

Key findings:

• No leak with PyArrow alone: When using PyArrow to read Parquet without converting to Pandas (i.e., no .to_pandas()), the memory leak does not occur.
• Leak with .to_pandas(): The memory leak appears during the conversion from Arrow to Pandas, suggesting the problem is tied to this process.
• No issue with Fastparquet or Polars: Fastparquet and Polars (even with PyArrow) do not exhibit this memory issue, reinforcing that the problem is in Pandas’ handling of Arrow data.

Reproduction Code

# dataset_creation.py

# just a fake dataset 

import pandas as pd
import numpy as np
import random
import string

np.random.seed(42)
random.seed(42)

def random_string(length):
    letters = string.ascii_letters
    return ''.join(random.choice(letters) for _ in range(length))

num_rows = 10**6  
col_types = {
    'col1': lambda: random_string(10),   
    'col2': lambda: np.random.randint(0, 1000),  
    'col3': lambda: np.random.random(),   
    'col4': lambda: random_string(5), 
    'col5': lambda: np.random.randint(1000, 10000),  
    'col6': lambda: np.random.uniform(0, 100),     
    'col7': lambda: random_string(8),   
    'col8': lambda: np.random.random() * 1000,  
    'col9': lambda: np.random.randint(0, 2),  
    'col10': lambda: random_string(1000)    
}

data = {col: [func() for _ in range(num_rows)] for col, func in col_types.items()}
df = pd.DataFrame(data)
df.to_parquet('random_dataset.parquet', index=True)

import os
file_size = os.path.getsize('random_dataset.parquet') / (1024**3) 
print(f"File size: {file_size:.2f} GB")

# memory_test.py

import pandas as pd 
import polars as pl
import gc
import pyarrow.parquet
import ctypes

data_path = 'random_dataset.parquet'

# To manually trigger memory release
malloc_trim = ctypes.CDLL("libc.so.6").malloc_trim

for _ in range(10): 
    df = pd.read_parquet(data_path, engine="pyarrow")
    # Also tested with:
    # df = pyarrow.parquet.read_pandas(data_path).to_pandas()
    # df = pl.read_parquet(data_path, use_pyarrow=True)
    
    del df  # Explicitly delete DataFrame
    
    for _ in range(3):  # Force garbage collection multiple times
        gc.collect()
        memory_info = psutil.virtual_memory()

    print(f"\n\nIteration number: {i}")
    print(f"Total Memory: {memory_info.total / (1024 ** 3):.2f} GB")
    print(f"Memory at disposal: {memory_info.available / (1024 ** 3):.2f} GB")
    print(f"Memory Used: {memory_info.used / (1024 ** 3):.2f} GB")
    print(f"Percentage of memory used: {memory_info.percent}%")

# Calling malloc_trim(0) is the only way we found to release the memory
malloc_trim(0)

Observations:

• Garbage Collection: Despite invoking the garbage collector multiple times, memory allocated to the Python process keeps increasing when .to_pandas() is used, indicating improper memory release during the conversion.
• Direct Use of PyArrow: When we import the data directly using PyArrow (without converting to Pandas), the memory usage remains stable, showing that the problem originates in the Arrow-to-Pandas conversion process.
• Manual Memory Release (ctypes): The only reliable way we have found to release the memory is by manually calling malloc_trim(0) via ctypes. However, we believe this is not a proper solution and that memory management should be handled internally by Pandas.

OS environment

Icon name: computer-vm
Chassis: vm
Virtualization: microsoft
Operating System: Red Hat Enterprise Linux 8.10 (Ootpa)
CPE OS Name: cpe:/o:redhat:enterprise_linux:8::baseos
Kernel: Linux 4.18.0-553.16.1.el8_10.x86_64
Architecture: x86-64

Conclusion

The issue seems to occur during the conversion from Arrow to Pandas, rather than being a problem within PyArrow itself. Given that memory is only released by manually invoking malloc_trim(0), we suspect there is a problem with how PyArrow handles memory management when converting the data to Panda. This issue does not arise when using Fastparquet engine or Polars instead of Pandas, further indicating that it is specific to the Pandas-Arrow interaction.

We recommend investigating how memory is allocated and released during the conversion from Arrow objects to Pandas DataFrames to resolve this issue.

Please let us know if further details are needed, and we are happy to assist.

Contributors:

• @Voltagabbana
• @okamiRvS
• @CarloNicolini
• @gfransvea

We would appreciate any feedback or insights from the maintainers and other contributors on how to improve memory management in this context.

Installed Versions

INSTALLED VERSIONS

commit : d9cdd2ee5a58015ef6f4d15c7226110c9aab8140
python : 3.10.14.final.0
python-bits : 64
OS : Linux
OS-release : 4.18.0-553.16.1.el8_10.x86_64
Version : #1 SMP Thu Aug 1 04:16:12 EDT 2024
machine : x86_64
processor : x86_64
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : en_US.UTF-8

pandas : 2.2.2
numpy : 2.0.0
pytz : 2024.1
dateutil : 2.9.0.post0
setuptools : 69.5.1
pip : 24.0
Cython : None
pytest : None
hypothesis : None
sphinx : None
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : 3.1.4
IPython : 8.26.0
pandas_datareader : None
adbc-driver-postgresql: None
adbc-driver-sqlite : None
bs4 : 4.12.3
bottleneck : None
dataframe-api-compat : None
fastparquet : 2024.5.0
fsspec : 2024.6.1
gcsfs : None
matplotlib : 3.9.0
numba : None
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : 17.0.0
pyreadstat : None
python-calamine : None
pyxlsb : None
s3fs : None
scipy : None
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
zstandard : None
tzdata : 2024.1
qtpy : None
pyqt5 : None

Component(s)

Parquet, Python