apache / apache/arrow

[Python] Dataset.to_batches accumulates memory usage and leaks

Open
#39,808 15 comments 25 reactions 1 assignee Claimed by @justinli500 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

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

If you want to read in a large parquet file or series of parquet files, the dataset reader accumulates the memory it allocates as you iterate through the batches.

To recreate:

```python
import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.dataset as ds

# Create dataset from large parquet file or files
dataset = ds.dataset('data/parquet', format='parquet')

# Iterate through batches
for batch in dataset.to_batches(batch_size=1000, batch_readahead=0, fragment_readahead=0):
print(pa.total_allocated_bytes())
pass
print(pa.total_allocated_bytes())
````

I am running this on OSX, which I believe uses `mimalloc` backend by default. It's worth noting that this is not the behavior that `ParquetFile.iter_batches` has. If you swap in that iterator, it will de-allocate the memory as soon as the batch leaves scope.

### Component(s)

Parquet, Python

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.