[BUG] Partitioned Parquet writes materialize a full extra device copy of the frame
- 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
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