[BUG] Partitioned Parquet write allocates a full extra copy of the frame; GPU idle ~70% on network FS

[VibhuJawa]

DRAFT for upstream rapidsai/cudf. Measured, not narrated. Two of the original production claims did not reproduce at single-node scale and are recorded as such below — please read that section before acting on this. Open before filing: the production configuration (100+ GPUs, max_file_size set) has not been reproduced end to end.

Describe the bug

Writing a partitioned Parquet dataset from cuDF costs exactly one extra full copy of the frame in device memory, on top of the write scratch a plain unpartitioned write already needs, and independent of filesystem. On a network filesystem it additionally leaves the GPU idle for most of the write, because it creates one file per partition per append.

Measurements

Workload: 500,000 rows × list<float32>[1024] plus an int32 partition key (one frame = 1.91 GiB, incompressible random data), partitioned by a 1,000-value centroid key, 4 appends. Single H100 80GB, cuDF 26.08.00a990 (nightly). Output ≈ 7.6 GiB (7,823 MiB) in every case.

"Extra device memory" below is memory allocated above the four already-resident frames — i.e. the writer's own working set, not a total peak.

Node-local disk (/tmp, ext3)

Path	Wall	Extra device mem	vs 1 frame	Files	GPU util (mean / median)	GPU idle
C unpartitioned baseline	7.3–7.4 s	6.42 GiB	3.36×	4	~20% / 3%	~1–2%
A to_parquet(partition_cols=)	8.2 s	8.34–8.35 GiB	4.37×	4,000	~12–13% / 2–6%	~29–32%
B ParquetDatasetWriter	8.2–8.3 s	8.34 GiB	4.37×	1,000	~11% / 2%	0%

Network filesystem (WekaFS)

Path	Wall	Extra device mem	vs 1 frame	Files	GPU util (mean / median)	GPU idle
C unpartitioned baseline	12.6–12.7 s	6.42 GiB	3.36×	4	~12% / 2%	~0–2%
A to_parquet(partition_cols=)	24.5–24.6 s	8.34–8.35 GiB	4.37×	4,000	~4% / 0%	69–74%
B ParquetDatasetWriter	12.0–12.5 s	8.34 GiB	4.37×	1,000	~12% / 4–5%	25–27%

Wall time and GPU utilization vary a few percent between runs; the memory figures are stable to the second decimal across every run we did, on both filesystems.

What the numbers say

1. The memory cost of partitioning is exactly one frame. Both partitioned paths allocate 8.34 GiB against the unpartitioned baseline's 6.42 GiB — a delta of +1.92 GiB against a 1.91 GiB frame. Identical on both filesystems and between the two partitioned paths. Partitioning materializes a complete second copy of the table.

2. The baseline is itself amplified, and that is a separate and larger problem. 6.42 GiB of scratch for a 1.91 GiB frame is 3.36× the frame in extra memory, i.e. a 4.36× total peak. That is not partition-specific: it is the leaf-element-count scaling filed as [BUG] Parquet write scratch scales with leaf ELEMENT count, not bytes: 4.4x amplification at fp32, 14.4x at int8 #12, which independently measures 4.36× for this dtype and predicts 6.42 GiB for this element count. The +1 frame here is what partitioning adds on top of that.

3. The wall-clock cost is a filesystem cost, not a GPU compute cost. Path A is ~1.12× the unpartitioned baseline on local disk but ~1.94× on WekaFS, and GPU idle time tracks it. The write is bound on small-file creation against a network filesystem. Producing 4,000 files where 4 would do is what makes that expensive.

Steps/Code to reproduce bug

Takes an output root as argv[1], so local and network filesystems can be compared. Data is cp.random.standard_normal, not zeros — all-zero embeddings compress to ~1 MiB and the write never performs representative I/O.

"""Partitioned Parquet write: extra device memory, file count, GPU idle time.
   python repro.py <output_root>"""
import os, statistics, sys, threading, time
import cudf, cupy as cp, pylibcudf as plc, pynvml
import rmm.statistics
from cudf.io.parquet import ParquetDatasetWriter

rmm.statistics.enable_statistics()
ROWS, DIM, CENTROIDS, APPENDS = 500_000, 1024, 1000, 4
ROOT, GIB = sys.argv[1], 2**30

class Mon:
    """Sample GPU utilization and free memory on a background thread."""
    def __init__(self):
        self.h = pynvml.nvmlDeviceGetHandleByIndex(0)
        self.stop, self.utils, self.min_free = threading.Event(), [], None
    def _run(self):
        while not self.stop.is_set():
            self.utils.append(pynvml.nvmlDeviceGetUtilizationRates(self.h).gpu)
            f = pynvml.nvmlDeviceGetMemoryInfo(self.h).free
            self.min_free = f if self.min_free is None else min(self.min_free, f)
            time.sleep(0.02)
    def __enter__(self):
        self.t = threading.Thread(target=self._run, daemon=True); self.t.start(); return self
    def __exit__(self, *a):
        self.stop.set(); self.t.join()

def make_df(seed):
    cp.random.seed(seed)
    arr = cp.random.standard_normal((ROWS, DIM), dtype=cp.float32)
    return cudf.DataFrame({
        "embedding": cudf.Series.from_pylibcudf(plc.Column.from_cuda_array_interface(arr)),
        "centroid": cudf.Series((cp.arange(ROWS) % CENTROIDS).astype("int32"))})

def run(label, fn):
    out = os.path.join(ROOT, label.split()[0])
    frames = [make_df(i) for i in range(APPENDS)]
    cp.cuda.Stream.null.synchronize()
    base_free = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0)).free
    rmm.statistics.push_statistics()
    with Mon() as m:
        t0 = time.perf_counter(); fn(frames, out)
        cp.cuda.Stream.null.synchronize(); wall = time.perf_counter() - t0
    rmm_extra = rmm.statistics.pop_statistics().peak_bytes
    nfiles, u = sum(len(f) for _, _, f in os.walk(out)), m.utils
    print(f"{label:<30} {wall:6.2f}s  extra: RMM {rmm_extra / GIB:5.2f} GiB / "
          f"NVML {(base_free - m.min_free) / GIB:5.2f} GiB  files {nfiles:>6,}  "
          f"GPU util mean {statistics.mean(u):4.1f}% median {statistics.median(u):4.1f}%  "
          f"idle {100 * sum(1 for x in u if x == 0) / len(u):4.1f}%", flush=True)
    del frames
    cp.get_default_memory_pool().free_all_blocks()

def path_c(frames, out):                     # baseline: unpartitioned
    os.makedirs(out, exist_ok=True)
    for i, df in enumerate(frames):
        df.to_parquet(os.path.join(out, f"p{i}.parquet"), index=False)

def path_a(frames, out):                     # to_parquet(partition_cols=)
    for df in frames:
        df.to_parquet(out, partition_cols=["centroid"], index=False)

def path_b(frames, out):                     # persistent ParquetDatasetWriter
    w = ParquetDatasetWriter(out, partition_cols=["centroid"], index=False)
    for df in frames:
        w.write_table(df)
    w.close()

pynvml.nvmlInit()
os.makedirs(ROOT, exist_ok=True)
print(f"cudf {cudf.__version__}  root={ROOT}")
print(f"one frame = {(ROWS * DIM * 4 + ROWS * 4) / GIB:.2f} GiB, {APPENDS} appends\n")
run("C  unpartitioned baseline", path_c)
run("A  to_parquet(partition_cols)", path_a)
run("B  ParquetDatasetWriter", path_b)

Output, node-local /tmp (ext3):

cudf 26.08.00a990  root=/tmp/repro7min
one frame = 1.91 GiB, 4 appends

C  unpartitioned baseline        7.39s  extra: RMM  6.42 GiB / NVML  6.43 GiB  files      4  GPU util mean 19.6% median  3.0%  idle  1.4%
A  to_parquet(partition_cols)    8.21s  extra: RMM  8.34 GiB / NVML  8.35 GiB  files  4,000  GPU util mean 12.2% median  6.0%  idle 28.9%
B  ParquetDatasetWriter          8.33s  extra: RMM  8.34 GiB / NVML  8.34 GiB  files  1,000  GPU util mean 11.7% median  2.0%  idle  0.0%

Same script, same node, output to WekaFS:

C  unpartitioned baseline       12.62s  extra: RMM  6.42 GiB / NVML  6.42 GiB  files      4  GPU util mean 12.4% median  2.0%  idle  1.6%
A  to_parquet(partition_cols)   24.51s  extra: RMM  8.34 GiB / NVML  8.35 GiB  files  4,000  GPU util mean  4.2% median  0.0%  idle 73.7%
B  ParquetDatasetWriter         12.45s  extra: RMM  8.34 GiB / NVML  8.34 GiB  files  1,000  GPU util mean 12.2% median  5.0%  idle 26.5%

Two production claims that did NOT reproduce

Stated plainly so this issue is not built on them:

• "ParquetDatasetWriter fixes the memory problem." It does not. Extra memory is identical to partition_cols= (8.34 GiB by RMM for both) on both filesystems. Whatever helped in production, it was not device memory.
• "ParquetDatasetWriter is slower than the one-shot path." At this scale it is faster or equal on both filesystems.

The most likely explanation for the divergence: production sets max_file_size, and that configuration is broken independently. It exhausts file descriptors and fails outright (#10), and per rapidsai/cudf#23378 it also mis-sizes list columns and over-rotates files. A path thrashing file rotation would plausibly present as "slower". We have not reproduced the production configuration end to end at 100+ GPU scale.

Production workload this models

Total output	~40 TB
Scale	100+ GPUs, one independent actor per GPU
Target filesystem	network filesystem (the measurements above used WekaFS)
Partition column	1,000–10,000 distinct values (k-means centroid id)
Payload	list<float32> embedding, up to 1024 values/row

for subgroup in [[file1, file2], [file3, file4, file5], ...]:
    df = cudf.read_parquet(subgroup)
    df["centroid"] = kmeans.predict(df["embedding"])
    df.to_parquet(output_path, partition_cols=["centroid"])   # path A

With 1,000 centroids, every subgroup contributes another file to every populated centroid directory, so file count grows as subgroups × centroids.

Expected behavior

• A partitioned write should not require a full extra copy of the frame in device memory. Ideally partitions are streamed out one at a time rather than materializing a fully partitioned duplicate.
• Writing N partitions should not cost N file creations per append when an appending writer is available.

Questions

1. Is the +1 frame allocation in the partition_cols= path inherent to the partition/split step, or is there a streaming formulation that avoids materializing the partitioned copy?
2. Why does ParquetDatasetWriter show the same extra allocation as the one-shot path? Naively it should be able to write partition-at-a-time.
3. Would user-controlled memory resources for intermediate allocations (rapidsai/cudf#20780, rapidsai/cudf#21282) let us bound this explicitly?
4. Is there a recommended pattern for "many partitions × many appends" — e.g. pre-sorting by the partition key so each write touches contiguous rows?

On measurement method

Memory above is scoped RMM statistics (rmm.statistics), which counts allocations directly. The script also prints an NVML free-memory delta sampled at 20 ms for comparison; on the default CudaMemoryResource the two agree closely.

Do not use NVML if an RMM pool is in play. With a 40 GiB PoolMemoryResource pre-reserved, NVML reports 0.537 GiB where RMM reports 8.990 GiB for a single-file write, and 0.006 / 0.000 GiB where RMM reports ~11.65 GiB for the partitioned paths — the pool is already reserved from the driver's point of view before measurement starts. NVML sees driver-visible memory, not allocations: it understates when the pool is pre-reserved and overstates when the pool grows inside the window. (Detail in #12.)

Environment overview

• Environment location: bare-metal, single node of a SLURM cluster
• Method of cuDF install: pip — cudf-cu12==26.8.0a990 (nightly wheel), rmm-cu12==26.8.0a62

Environment details

• NVIDIA H100 80GB HBM3, CUDA 12.9, Python 3.12, cupy 14.1.1
• Filesystems compared: node-local ext3 (/tmp) and WekaFS

Additional context

• rapidsai/cudf#23378 — ParquetDatasetWriter max_file_size overestimates list columns and creates too many files. Same workload, same writer; that issue covers file count, this one covers memory and throughput.
• rapidsai/cudf#20443 (filesystem support in to_parquet), rapidsai/cudf#22367 (higher parquet read memory 26.04 vs 26.02).

Related drafts from us (this fork, to be filed upstream alongside): #12 (write scratch scales with leaf element count — the amplified baseline this issue measures a delta against) · #10 (max_file_size exhausts file descriptors).