NVIDIA / NVIDIA/cudf

[BUG] Partitioned Parquet writes materialize a full extra device copy of the frame

Open
#23,502 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

## Describe the bug

A partitioned Parquet write allocates a complete additional device-sized table before encoding. For a 500,000-row `list[1024]` frame, scoped RMM statistics show an additional **1.92 GiB**, essentially one full **1.91 GiB** frame, on top of the unpartitioned writer working set.

This is independent of the filesystem. On WekaFS the many-small-files path also leaves the GPU idle for most of the write.

## Reproduction and measurements

The full RMM/NVML reproducer is in [VibhuJawa/cudf#7](https://github.com/VibhuJawa/cudf/issues/7). Workload: 4 already-resident frames, 500,000 rows each, `list[1024]` plus a 1,000-value int32 partition key, 4 appends, cuDF `26.08.00a990`, H100 80GB.

| Path | Extra device memory | Files | WekaFS GPU idle |
|---|---:|---:|---:|
| unpartitioned baseline | 6.42 GiB | 4 | ~0–2% |
| `to_parquet(partition_cols=...)` | 8.34 GiB | 4,000 | ~69–74% |
| `ParquetDatasetWriter` | 8.34 GiB | 1,000 | ~25–27% |

The partitioned delta is **+1.92 GiB**, equal to one input frame. The same memory figures were measured on node-local ext3. These are scoped RMM allocation counters; the input used incompressible random values.

## Source indication

Both [`_get_partitioned`](https://github.com/rapidsai/cudf/blob/main/python/cudf/cudf/io/parquet.py#L1711-L1746) and [`ParquetDatasetWriter.write_table`](https://github.com/rapidsai/cudf/blob/main/python/cudf/cudf/io/parquet.py#L2110-L2170) call [`_get_groups_and_offsets`](https://github.com/rapidsai/cudf/blob/main/python/cudf/cudf/io/parquet.py#L1749-L1772), which performs a groupby and returns a materialized `grouped_df` before passing it to the Parquet writer.

## Expected behavior

Partitioned output should stream partitions or otherwise avoid retaining a full second copy of the frame when memory pressure matters. If materialization is intentional, the peak-memory cost should be documented and a bounded/streaming option provided.

The small-file/network-filesystem throughput effect is a related but separate concern from the extra device allocation.

Contributor guide

Open the contributing guide

Research direction

Start in python/cudf/cudf/io/parquet.py by reading _get_partitioned, ParquetDatasetWriter.write_table, and _get_groups_and_offsets. Run the RMM/NVML reproducer in VibhuJawa/cudf#7 and compare partitioned and unpartitioned writes. Done means partitioned output avoids retaining a full extra device-sized frame, or clearly documents the bounded/streaming behavior and its peak-memory cost.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.