apache / apache/arrow

[Python] OOM with Dataset.scanner but ok with ParquetFile.iter_batches

Open
#44,799 4 comments 0 reactions 0 assignees View on GitHub
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.

#### Setup

I am using `pyarrow` version `18.0.0`.

I am running my tests on an AWS `r6g.large` instance running Amazon Linux. (I also attempted using instances with larger memory in case the problem was that there was some base-level memory needed irrespective of minimal batch sizes and readahead, but this didn't help.)

My data consists of parquet files in S3, varying in size from a few hundred kB to ~ 1GB, for a total of 3.4GB. This is a sample subset of my actual dataset which is ~ 50GB.

#### Problem description

I have a set of parquet files with very small row-groups, and I am attempting to use the `pyarrow.dataset` API to transform this into a set of files with larger row-groups. My basic approach is `dataset -> scanner -> write_dataset`. After running into OOM problems with default parameters, I ratcheted down the read and write batch sizes and concurrent readahead:
``` python
from pyarrow import dataset as ds

data = ds.dataset(INPATH, format='parquet')

# note the small batch size and minimal values for readahead
scanner = data.scanner(
batch_size=50,
batch_readahead=1,
fragment_readahead=1
)

# again, note extremely small values for output batch sizes
ds.write_dataset(
scanner,
base_dir=str(OUTPATH),
format='parquet',
min_rows_per_group=1000,
max_rows_per_group=1000
)
```
Running this results in increasing memory consumption (monitored using `top`) until the process maxes out available memory and is finally killed.

What worked to keep memory use under control was to replace the `dataset` scanner with `ParquetFile.iter_batches` as below:

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

def batcherator(filepath, batch_size):
for f in filepath.glob('*.parquet'):
with pq.ParquetFile(f) as pf:
yield from pf.iter_batches(batch_size=batch_size)

scanner = batcherator(INPATH, 2000) # it's fine with higher batch size than previous

ds.write_dataset(
scanner,
base_dir=str(OUTPATH),
format='parquet',
min_rows_per_group=10_000, # again, higher values of write batch sizes
max_rows_per_group=10_000
)
```

Since nothing's really changing on the `dataset.write_dataset` side, it seems like there's some issue with runaway memory use on the `scanner` side of things?

The closest I could find online was this DuckDB issue https://github.com/duckdb/duckdb/issues/7856 which in turn pointed to this arrow issue https://github.com/apache/arrow/issues/31486 but this seems to hint more at a problem with `write_dataset`, which for me seemed ok once I replaced how I am reading in the data.

### Component(s)

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.