apache / apache/arrow

[Python] Small max_rows_per_group causes high memory usage while iterating over each fragment's row_groups (and never released)

Open
#43,387 0 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.

Let's say we use the following code to create some test parquet files
```python

import pandas as pd
from itertools import repeat
from collections import defaultdict

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

schema: pa.schema = pa.schema(
[
("id", pa.int64()),
("tag", pa.string()),
]
)

fast_write_options: dict = dict(
compression="lz4",
version="2.6",
write_statistics=True,
)

for i, suffix in enumerate({"", "_10", "_20", "_30", "_40", "_50", "_60", "_70", "_80", "_90", "_100"}):
outputs = []

for tag in {"first", "second", "third", "fourth", "fifth", "sixth", "seventh", "eighth", "nineth", "tenth"}:
outputs.append(
pd.DataFrame({"id": list(range((i + 1) * int(1e6))), "tag": list(repeat(tag, (i + 1) * int(1e6)))})
)

df = pd.concat(outputs, ignore_index=True)

writer = pq.ParquetWriter(
str(f"./test{suffix}.parquet"),
schema=schema,
filesystem=pa.fs.LocalFileSystem(),
**fast_write_options,
)

table = pa.Table.from_pandas(df, schema=schema, preserve_index=False)
writer.write_table(table)

writer.close()
```

Then, I partition them based on the `tag` field

```python
dataset_files = []
for f in {"", "_10", "_20", "_30", "_40", "_50", "_60", "_70", "_80", "_90", "_100"}:
dataset_files.append(f"./test{f}.parquet")

input_datasets = ds.dataset(
dataset_files,
schema=schema,
format="parquet",
filesystem=pa.fs.LocalFileSystem(),
)

scanner = input_datasets.scanner(
batch_readahead=4,
fragment_readahead=1,
use_threads=True,
)
ds.write_dataset(
scanner,
base_dir="base_part_dir",
basename_template=f"part-{{i}}.parquet",
format="parquet",
filesystem=pa.fs.LocalFileSystem(),
partitioning=ds.partitioning(pa.schema([("tag", pa.string())]), flavor="hive"),
max_rows_per_group=1_024,
min_rows_per_group=1_024,
max_rows_per_file=20 * 10_480,
existing_data_behavior="overwrite_or_ignore",
file_options=ds.ParquetFileFormat().make_write_options(**fast_write_options),
use_threads=False,
)
```

And, try to read the row_groups of each fragment

```python
import time

import pyarrow as pa
import pyarrow.dataset as ds

schema: pa.schema = pa.schema(
[
("id", pa.int64()),
("tag", pa.string()),
]
)

def main():
intermediate_dataset = ds.dataset(
"base_part_dir",
schema=schema,
format="parquet",
partitioning=ds.partitioning(pa.schema([("tag", pa.string())]), flavor="hive"),
)

keys = []
for f in intermediate_dataset.get_fragments():
keys.append(ds.get_partition_keys(f.partition_expression)["tag"])
f.row_groups

if __name__ == "__main__":
time.sleep(10)

print("started")
while True:
main()

print("Sleepig...")
time.sleep(10)
```

On the initial execution, it uses approximately 128MB of memory and, after having executed the `main()` once, it goes up to 900MB, and it stays there after successive iterations.

Changing the values of **max_rows_per_group / min_rows_per_group** to 65_536, the memory usage does not increase at all, or very slightly.

Therefore, my question is, why is this happening? Firstly, I would expect that reading metadata, i.e., `f.row_groups`, would cause minimal memory usage. And secondly, I would expect that, after the execution of `main()` is finished, the memory should get released.

Is this a bug, anticipated or I use the library in the wrong way? Is there something that I can do to mitigate this?

### 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.