Support Loading Parquet Files in WholeGraph by alexbarghi-nv · Pull Request #497 · rapidsai/cugraph-gnn

alexbarghi-nv · 2026-07-21T21:20:48Z

Summary

Adds direct Parquet input support for WholeMemory tensors and embeddings, including the cuGraph-PyG DistTensor and DistEmbedding file-loading APIs.

The implementation is designed for datasets that may be several terabytes: it allocates the final WholeMemory destination first, then decodes and copies only bounded Parquet batches into each rank's local partition. It never loads or concatenates the complete source dataset in host memory.

Implementation

supports binary, parquet, and auto file formats
uses Parquet footer metadata to infer and optionally validate tensor shape before allocation
makes expected_shape optional for file input; it is used only for validation and one-column dimensionality disambiguation
infers last_dim_size for Parquet input and requires it for raw binary input
treats one-column Parquet input as a 1-D (N,) tensor by default; expected_shape=(N, 1) requests the 2-D representation
keeps the destination dtype explicit and warns when Parquet column types differ before converting bounded batches
adds fail_on_dtype_mismatch=True to turn that warning into a metadata-validation error before allocation
treats a file list as one logical row-major tensor
maps each rank's contiguous WholeMemory row interval onto the input files and reads only intersecting row groups and rows
converts one bounded columnar batch at a time into the dense row-major layout required by WholeMemory
writes batches directly into the final CPU or CUDA WholeMemory allocation
preserves the native binary loader as the lowest-overhead path
supports explicit tensor_entry_partition / partition_book values; otherwise WholeMemory applies its default partitioning
keeps round-robin loading binary-only

Why direct PyTorch and NumPy file input are not supported

PyTorch's public file-loading API does not provide a bounded, windowed read interface. torch.load(..., mmap=True) avoids eagerly copying the complete tensor, but pages touched during a sequential scan can remain resident while the full-file mapping is open. Keeping RSS strictly bounded would therefore depend on platform-specific page-eviction calls, while older PyTorch serialization formats may not support mmap at all.

That behavior is not a good fit for a loader whose primary requirement is predictable host-memory use at multi-terabyte scale. Callers should use Parquet for structured input or convert data to the native binary format for the lowest possible overhead. .pt and .pth files are explicitly rejected when format detection is used.

The previous .npy path had the same issue: np.load(..., mmap_mode="c") avoided an eager copy, but pages touched while copying a rank's partition could remain resident and make peak RSS scale with that partition. Direct .npy loading has therefore also been removed. Known .pt, .pth, and .npy extensions are rejected even when callers attempt to force the binary format, preventing serialized headers from being misinterpreted as raw tensor data.

Memory behavior and tests

The tests cover:

Parquet correctness for CPU and CUDA WholeMemory destinations
Parquet integration through cuGraph-PyG DistTensor and DistEmbedding
metadata-only row-count and schema validation
inferred multi-column shape and optional expected-shape validation
one-column (N,) and (N, 1) behavior
warnings for Parquet/destination dtype mismatches
strict dtype mismatch rejection through the low-level and cuGraph-PyG APIs
explicit rejection of direct .pt, .pth, and .npy input
requested row ranges and multi-row-group input
non-numeric column and shape validation
end-to-end peak host RSS for CPU and CUDA destinations
reader-only peak host RSS
memory scaling from 32 MiB to 128 MiB datasets in fresh subprocesses

The scaling tests subtract only the required CPU WholeMemory destination allocation and verify that the remaining reader/conversion overhead stays within a fixed bound and does not grow materially when the dataset size increases 4x.

Local validation

pylibwholegraph Parquet suite: 21 passed
cuGraph-PyG Parquet from_file integration: 4 passed
pre-commit hooks: passed
git diff --check: passed

copy-pr-bot · 2026-07-21T21:20:51Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

copy-pr-bot · 2026-07-23T20:14:05Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-07-25T22:08:16Z

Greptile Summary

Adds bounded, direct Parquet loading for WholeMemory-backed tensors and embeddings.

Infers and validates Parquet shape and dtype metadata before allocation.
Streams intersecting row batches into each rank’s local CPU or CUDA partition.
Exposes binary, Parquet, and automatic format selection through cuGraph-PyG APIs.
Adds PyArrow packaging dependencies and correctness, integration, and memory-scaling tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
python/pylibwholegraph/pylibwholegraph/torch/tensor.py	Implements format detection, Parquet metadata validation, bounded batch conversion, and rank-local file loading.
python/pylibwholegraph/pylibwholegraph/torch/embedding.py	Extends embedding construction to infer Parquet dimensions and stream data into the allocated embedding.
python/cugraph-pyg/cugraph_pyg/tensor/dist_tensor.py	Exposes the new file-format, expected-shape, dimension, and strict dtype options through DistTensor and DistEmbedding.
python/cugraph-pyg/cugraph_pyg/tensor/utils.py	Forwards file metadata and format options from cuGraph-PyG into the WholeMemory constructors.
dependencies.yaml	Adds PyArrow to pylibwholegraph runtime and test dependency sets.

Sequence Diagram

sequenceDiagram
    participant API as DistTensor / DistEmbedding
    participant Meta as Parquet metadata reader
    participant WM as WholeMemory allocator
    participant Reader as Bounded Parquet reader
    participant Rank as Rank-local partition
    API->>Meta: Inspect schemas and row counts
    Meta-->>API: Inferred or validated shape
    API->>WM: Allocate final distributed tensor
    WM-->>API: Rank-local row interval
    API->>Reader: Read intersecting row groups in batches
    Reader-->>Rank: Convert and copy dense row-major batches
    Rank->>WM: Synchronize completed load

_{Reviews (4): Last reviewed commit: "improve error message" | Re-trigger Greptile}

alexbarghi-nv · 2026-07-27T21:33:57Z

/ok to test 59185e5

alexbarghi-nv · 2026-07-27T23:09:24Z

/ok to test 0aa7e8e

linhu-nv

This looks good to me.

jameslamb

Reviewed for packaging-codeowners, left some very very minor feedback. Do what you want with it, including ignoring it or addressing in a follow-up. Marking this "approve" so you have what you need to merge without another review.

I only lightly skimmed the Python code, deferring to the other reviews on correctness, performance, etc.

jameslamb · 2026-07-29T21:52:04Z

    common:
      - output_types: [conda, pyproject, requirements]
        packages:
+          - *pyarrow


Suggested change

- *pyarrow

Since pyarrow is now a hard required dependency, I'd keep it out of the [test] extra. I think we generally prefer to only have that extra provide additional packages on top of what's in dependencies or to provide different constraints.

For example, cudf declares pyarrow in its runtime dependencies but not its [test] extra: https://github.com/rapidsai/cudf/blob/cfb5a7d8ce9c902844e216054a49b36daebc5491/python/cudf/pyproject.toml#L34

jameslamb · 2026-07-29T22:03:51Z

 dependencies = [
    "libwholegraph==26.10.*,>=0.0.0a0",
    "numpy>=2.0,<3.0",
+    "pyarrow",


Just to write this somewhere... I think it's totally fine to start with this unconstrained and only add constraints as you discover issues, especially with the small slice of the pyarrow API this PR is using.

Actually now that I think about it - do we have a RAPIDS-wide pyarrow constraint?

Not that we should apply it here but just as something that's nice to know

Good question. I think the answer is "not really, everything just kind of follows what pandas / cuDF do".

Looks like cudf today is the only one with an explicit pin on it: https://github.com/search?q=pyarrow+path%3Adependencies.yaml+%28org%3Arapidsai+OR+org%3ANVIDIA%29+AND+NOT+is%3Aarchived&type=code

Co-authored-by: James Lamb <jaylamb20@gmail.com>

alexbarghi-nv · 2026-07-29T22:08:48Z

/ok to test f7bafc5

alexbarghi-nv · 2026-07-29T22:09:57Z

/merge

initial write

1e99e0c

alexbarghi-nv added 4 commits July 22, 2026 18:35

initial write

8802d19

Merge remote-tracking branch 'upstream/main' into wg-load-parquet

ba1fc1a

update to support low-level API

62b8469

merge and update

8cdcd14

alexbarghi-nv added breaking Introduces a breaking change feature request New feature or request labels Jul 23, 2026

alexbarghi-nv self-assigned this Jul 23, 2026

update tests to include CPU, fix copy issue, add clarifying comments

624953b

alexbarghi-nv added 5 commits July 24, 2026 10:18

more test improvements'

a589e4b

Merge remote-tracking branch 'upstream/main' into wg-load-parquet

68bbeb4

remove pyt loading, improve test coverage

eb8af73

simplify api, add option to fail on dtype mismatch

11d44d9

drop numpy support

28fdabc

alexbarghi-nv marked this pull request as ready for review July 25, 2026 22:06

alexbarghi-nv requested review from a team as code owners July 25, 2026 22:06

alexbarghi-nv requested a review from bdice July 25, 2026 22:06

resolve merge conflict

59185e5

relax memory limit, remove pandas usage

0aa7e8e

linhu-nv approved these changes Jul 28, 2026

View reviewed changes

alexbarghi-nv mentioned this pull request Jul 28, 2026

Support DistMatrix Creation from Parquet Files #509

Draft

tingyu66 approved these changes Jul 29, 2026

View reviewed changes

jameslamb approved these changes Jul 29, 2026

View reviewed changes

jameslamb removed the request for review from bdice July 29, 2026 22:06

improve error message

f7bafc5

Co-authored-by: James Lamb <jaylamb20@gmail.com>

Uh oh!

Conversation

alexbarghi-nv commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

Why direct PyTorch and NumPy file input are not supported

Memory behavior and tests

Local validation

Uh oh!

copy-pr-bot Bot commented Jul 21, 2026

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

greptile-apps Bot commented Jul 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

alexbarghi-nv commented Jul 27, 2026

Uh oh!

alexbarghi-nv commented Jul 27, 2026

Uh oh!

linhu-nv left a comment

Choose a reason for hiding this comment

Uh oh!

jameslamb left a comment

Choose a reason for hiding this comment

Uh oh!

jameslamb Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jameslamb Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

alexbarghi-nv Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

alexbarghi-nv Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

jameslamb Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

alexbarghi-nv commented Jul 29, 2026

Uh oh!

alexbarghi-nv commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

alexbarghi-nv commented Jul 21, 2026 •

edited

Loading

greptile-apps Bot commented Jul 25, 2026 •

edited

Loading