[FEA] Improving read parquet + udf + partitioned parquet write at scale
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
**Is your feature request related to a problem? Please describe.**
**TLDR: Provide a managed streaming Parquet reader and partitioned writer that maximize throughput while sharing GPU memory safely with application compute. Users supply files, transformations, and partition columns; cuDF owns I/O execution and resource management.**
Imagine thousands of Parquet files containing terabytes of data with a `list[float]` column, such as a numeric vector per row. A user wants to read the data, perform a UDF, and write the results partitioned by a column.
The approximately 2-billion-element limit applies to the list's child column. With 1,024 values per row, a DataFrame can accommodate approximately 2.1 million rows: about **4 GiB of FP16 payload or 8 GiB of FP32 payload**, before offsets, metadata, and temporary allocations. Partitioned writing adds grouping and encoding/compression allocations, whose sizes depend on the workload. Even with these allocations, a single batch **can leave considerable headroom on an 80 GB GPU**. Using that headroom for concurrent batches could improve throughput, but currently leaves applications responsible for batching, memory estimates, and I/O concurrency.
Curator attempted this in https://github.com/NVIDIA-NeMo/Curator/pull/2396, but the team decided against merging the additional application-side complexity. We would like to see if cuDF can provide some reusable I/O capabilities needed for this workload.
### What we tried in Curator
The read → UDF → write pass used **up to four threads on one GPU**, with each thread performing read → UDF → partitioned write sequentially for its assigned batches. Different threads could overlap these operations. Each thread owned its UDF instance and persistent partition writer; shared application state used by the UDF was read-only.
Simplified pseudocode of that approach (helper names describe behavior, not proposed public APIs):
```python
# Configure before cuDF is imported in each worker process.
CUDF_PER_THREAD_STREAM = "1"
KVIKIO_NTHREADS = "16" # Four I/O threads per requested application worker.
# Keep Curator's existing KVIKIO_AUTO_DIRECT_IO_WRITE=0 for SSD/Lustre.
# Initialize shared read-only UDF state and release setup buffers.
shared_state = initialize_shared_udf_state()
budget = int(free_gpu_memory() * 0.9)
groups, worker_count = plan_file_groups(
files,
max_workers=4,
memory_budget=budget,
# Respect both the child-element limit and estimated peak memory:
# decoded input + UDF scratch + partition/write scratch + metadata.
)
wait_for_shared_state_initialization_on_existing_streams()
def process_groups(worker_id, groups):
# Select the GPU and per-thread default stream in this new thread.
# Route CuPy allocations through rmm_cupy_allocator as well as cuDF's RMM.
with worker_gpu_context(device, cupy_allocator=rmm_cupy_allocator):
udf = make_worker_udf(shared_state)
with persistent_partition_writer(
output_path, worker_id, partition_cols=["partition_column"]
) as writer:
for paths in groups:
frame = cudf.read_parquet(paths)
output = udf(frame) # Includes the partition_column used for writing.
writer.write(output) # Synchronous per worker; overlaps other workers.
del frame, output
# One shared rmm.mr.PoolMemoryResource wrapping the existing device resource.
# Install before starting threads; restore the previous resource after they finish.
# Starts at up to 1 GiB and grows to the budget, rather than reserving it all upfront.
with shared_rmm_pool(initial_size=min(GiB, budget), maximum_size=budget):
with ThreadPoolExecutor(max_workers=worker_count) as executor:
futures = [
executor.submit(process_groups, i, groups[i::worker_count])
for i in range(worker_count)
]
for future in futures:
future.result()
```
The application also had to implement the following:
- Estimate concurrent working sets, including temporary UDF buffers and retained metadata such as strings. The RMM pool bounds allocations routed through it, not every allocation made by application/library code.
- Maintain a writer per thread: use `cupy.bincount` for row counts per partition (our partition keys were nonnegative integer IDs), discover new partitions with `numpy.flatnonzero`, sort/group rows, compute offsets, and reuse `ParquetWriter` sinks with `partitions_info`.
- Rotate files before accumulated list-child element counts exceeded the limit, close and discard old writers, use distinct filenames across threads, and preserve the general I/O path for storage options/custom write arguments.
These are implementation details we would prefer applications not to own. The exact thread counts, 90% heuristic, and allocator policy above are what we experimented with, not requirements for an upstream design.
A separate cuDF **26.10.0a517** write-only comparison used four threads on 8.2 million rows with 384 values per list and 1,000 partitions. Native `ParquetDatasetWriter(max_file_size=...)` took approximately **25 seconds**, versus **3–3.5 seconds** for the custom writer. The nested-size correctness fix ([#23378](https://github.com/NVIDIA/cudf/issues/23378), fixed by [#23907](https://github.com/NVIDIA/cudf/pull/23907)) was present, but per-partition size estimation remained expensive. A separate forced-rotation test also showed output descriptors accumulating while the dataset writer remained alive, including after `close()`. These observations motivate efficient native size accounting and bounded writer resources, in addition to concurrency.
### Sample API
Illustrative proposed API:
```python
with cudf.ParquetDatasetWriter(
output_path,
partition_cols=["partition_column"],
storage_options=storage_options,
) as writer:
for batch in cudf.ParquetDatasetReader(input_files):
output = udf(batch) # Returns a DataFrame including partition_column.
writer.write(output)
```
The reader and writer should coordinate batching, prefetch, and outstanding writes automatically, without normally requiring application memory budgets, thread counts, or environment variables. Each batch should respect the column limit. `write()` may return before completion, with managed buffer lifetimes, CUDA dependencies, and backpressure; context exit should wait for completion and propagate errors. Please document how this coexists with application UDF allocations.
### Additional context
A downstream application stage reads each partition separately, so the output must retain the directory layout defined by `partition_cols`. We need string metadata and repeated appends to work on both SSD and Lustre with `KVIKIO_AUTO_DIRECT_IO_WRITE=0`. The goal is higher throughput with bounded live memory and file descriptors as dataset size grows.
Related issues: [list child-element limit #23493](https://github.com/NVIDIA/cudf/issues/23493), [extra partitioning copy #23502](https://github.com/NVIDIA/cudf/issues/23502), [writer descriptor retention #23501](https://github.com/NVIDIA/cudf/issues/23501), and [KvikIO configuration being overwritten #23786](https://github.com/NVIDIA/cudf/issues/23786).
Contributor guide
Research direction
Start by reviewing the existing cudf.read_parquet entry point and ParquetDatasetWriter APIs, then read related issues #23493, #23502, #23501, and #23786. The requested outcome is a managed streaming reader and partitioned writer with coordinated batching, prefetch, backpressure, bounded memory and file descriptors, and correct completion and error propagation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend-api-design, data-engineering, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100