apache / apache/arrow

[Python][C++] How to limit the memory consumption of to_batches()

Open
#33,759 16 comments 0 reactions 0 assignees View on GitHub
Component: Python Type: usage
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 18h
Merged PRs (30d)
91

Description

In order to get the unique values of a column of a 500GB Parquet dataset (made of 13 000 fragments) on a computer with 12GB of memory, I chose to use to_batches() as following :

`import pyarrow as pa`
`import pyarrow.dataset as ds`

`partitioning = ds.partitioning( pa.schema([(timestamp, pa.timestamp("us"))]),flavor="hive",)`
`unique_values = set()`
`dataset = ds.dataset(path, format="parquet", partitioning=partitioning)`
`batch_it = dataset to_batches(columns=[column_name])`
`for batch in batch_it:`
` unique_values.update(batch.column(column_name).unique())`

The problem is that the process quickly accumulates memory and exceeds the amount available.
When I put a breakpoint on the line "for batch in batch_it", the process continues to accumulate memory until it crashes.

I understand that to_batches readahead but I thought I could limit it with "fragment_readahead" parameter. Is there a way to limit readahead ? Is there a way to "free" memory after a batch has been consumed ?
Is there another way to go ? My first try was using to_table() but it needs 20GB of memory in that case. It seems that to_batches would also need 20GB

### Component(s)

Python

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the Python dataset.to_batches() example against the reported Parquet workload, focusing on fragment_readahead and memory growth while batches are consumed. Compare the observed behavior with the intended readahead and memory-lifetime behavior; done means the supported way to limit memory is established or the issue has a focused fix with regression coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
data-engineering, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.