Skip to content

[Python] Dataset.to_batches accumulates memory usage and leaks #39808

Description

@akoumjian

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:

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

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions