apache / apache/arrow

[Python][Parquet] Memory Leak when using ParquetWriter

Open
#45,971 1 comment 0 reactions 0 assignees View on GitHub
Component: Parquet 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.

We get data batches from BigQuery and write them to parquet. Parquet writer eats up all memory and crashed the pod.

It's no JMalloc or whatever since we do not see the issue when we periodically create new ParquetWriter instances.

This code leaks:

```py
from google.cloud import bigquery
from google.cloud.bigquery import _pandas_helpers
from pyarrow import parquet

client = bigquery.Client(project=...)
job = client.get_job(job_id=...)

result = job.result()
arrow_schema = _pandas_helpers.bq_to_arrow_schema(result.schema)
bqstorage_client = client._ensure_bqstorage_client()
with parquet.ParquetWriter(where="result.parquet", schema=arrow_schema) as writer:
for batch in result.to_arrow_iterable(
bqstorage_client=bqstorage_client,
max_queue_size=1,
max_stream_count=1,
):
writer.write_batch(batch)
```

Initially we though that the bug was in BigQuery, but we were wrong. https://github.com/googleapis/python-bigquery/issues/2151

Proof:

Changing from
```
with parquet.ParquetWriter(where="result.parquet", schema=arrow_schema) as writer:
for batch in result.to_arrow_iterable(...):
writer.write_batch(batch)
```

to
```
for batch in result.to_arrow_iterable(...):
with parquet.ParquetWriter(where="result.parquet", schema=arrow_schema) as writer:
writer.write_batch(batch)
```
makes the memory leak go away. Which means that the leak was in `ParquetWriter`.

Versions: "pyarrow==19.0.1", "pyarrow==16.1.0"

### Component(s)

Parquet

![Image](https://github.com/user-attachments/assets/8df18d91-ec61-43d3-93c5-328a9e2f5826)

Update: This code speedruns OOMKilled without any BigQuery:

```
from pyarrow import dataset as ds

data = ds.dataset(input_uri, format="parquet")
ds.write_dataset(
data,
base_dir="/tmp/merge_parquet_files/",
format="parquet",
)
```

Update 2:

Setting `write_statistics=False` does not help.

```
parquet.ParquetWriter(where=output_path, schema=arrow_schema, write_statistics=False) as writer
```

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.