[Python][Parquet] 21.0.0 release introduced memory leak when reading parquet

[Tom-Newton]

Describe the bug, including details regarding any error messages, version, and platform.

Since 21.0.0 memory usage grows a lot when repeatedly reading a parquet dataset on the local disk. With version 20.0.0 the memory usage increased much less.

Script to reproduce

import tempfile
import numpy
import pyarrow
import pyarrow.dataset
import pyarrow.parquet
from memory_profiler import profile

def test_memory_leak():
    num_columns = 10
    num_rows = 5_000_000

    data = {f"col_{i}": numpy.random.rand(num_rows) for i in range(num_columns)}
    table = pyarrow.Table.from_pydict(data)

    with tempfile.TemporaryDirectory() as temp_dir:
        pyarrow.dataset.write_dataset(table, temp_dir, format="parquet")

        @profile
        def read():
            return pyarrow.dataset.dataset(temp_dir).to_table()

        for _ in range(50):
            read()

if __name__ == "__main__":
    test_memory_leak()

Environment

Ubuntu 24.04.2 LTS
Tested python 3.10.15 and python 3.12.3

Python packages:

$ pip freeze
memory-profiler==0.61.0
numpy==2.3.2
psutil==7.0.0
pyarrow==21.0.0

When using pyarrow==21.0.0 the memory usage increases with the iterations. After the first read its at about 1.5GiB. After the 50th read its at about 20GiB. If I run the same test with pyarrow==20.0.0 the memory usage still increases slightly with the iterations but its still less than 2GiB after the 50th iteration.

Debugging

I ran a git bisect and identified #45979 as the change point. Building from the 21.0.0 release commit with ARROW_MIMALLOC=OFF also solves the problem.

Component(s)

C++, Python

[pitrou]

Ok, two questions:

1. how do you see the leak? are you measuring RSS for example? (Edit: ok, you're using memory_profiler which AFAIU measures RSS by default)
2. do you still notice a leak if you regularly run pa.default_memory_pool().release_unused()? see https://arrow.apache.org/docs/python/generated/pyarrow.MemoryPool.html#pyarrow.MemoryPool.release_unused

[pitrou]

Also note that you don't need to rebuild PyArrow, you can just use the ARROW_DEFAULT_MEMORY_POOL environment variable to switch to another memory pool.

[lucasKoeb]

Experienced similar issues after updating to 21.0.0 and had to revert back to 20.0.0 for now.

Could reproduce the issue using the provided sample script on python 3.12.10 on linux and windows.
Tried some variations on the reproduction script:

Using the ARROW_DEFAULT_MEMORY_POOL environment variable

Using the environment variable to change the default to system or jemalloc on linux, and also changing it to system on windows worked. On these scenarios, memory usage as reported by memory_profiler is kept under 2GiB.

Using pyarrow.set_memory_pool()

Using the set_memory_pool method to change the default memory pool does not prevent the memory leak, resulting in the same increases and peak memory usage as mimalloc.

Even though the call seems to work, and pyarrow.default_memory_pool() correctly reports the memory pool I have configured, this seems to have no practical effect on the allocations.

Tested this by setting the pool to pyarrow.system_memory_pool() and pyarrow.jemalloc_memory_pool() on linux and to pyarrow.system_memory_pool() on windows.

Passing a specific memory pool to the to_table() method is also ineffective.

Calling pa.default_memory_pool().release_unused() in between iterations

Changed the script to call release_unused after each read, tried on linux and windows. This does not prevent the leak, and the peak memory usage remains the same.

        for _ in range(50):
            read()
            pyarrow.default_memory_pool().release_unused()

[pitrou]

I will temptatively ping @daanx in case he can suggest how to diagnose this mimalloc memory consumption issue.

[pitrou]

Ok, I've tried the reproducer locally. It seems that switching to mimalloc 3.0.9 and 3.1.5 (as in #47589) would entirely fix the issue.

[pitrou]

Using the set_memory_pool method to change the default memory pool does not prevent the memory leak, resulting in the same increases and peak memory usage as mimalloc.

I've opened #47595 for this.