Support DistMatrix Creation from Parquet Files

[alexbarghi-nv]

Summary

Follow-up to #497

Extends direct Parquet loading to support column projection and memory-efficient COO DistMatrix construction.

A two-column Parquet edge list can now be loaded as two independent 1-D DistTensor objects:

matrix = DistMatrix.from_file(
    "edges.parquet",
    source="src",
    destination="dst",
    dtype=torch.int64,
    device="cpu",
)

Ordered file lists are also supported:

matrix = DistMatrix.from_files(
    parquet_files,
    source="src",
    destination="dst",
    dtype=torch.int64,
)

Changes

• Adds a columns option to the pylibwholegraph Parquet loader:
  • columns="src" produces a 1-D (E,) tensor.
  • columns=["a0", "a1"] produces an (E, 2) tensor in the requested order.
  • columns=None retains the existing physical-column-order behavior.
• Exposes columns through DistTensor.
• Adds DistMatrix.from_file() and DistMatrix.from_files().
• Adds source and destination column selection, defaulting to src and dst.
• Supports ordered multi-file datasets, device, backend, partition book, expected shape, and strict dtype validation.
• Validates missing, duplicate, empty, and nonnumeric column selections.
• Rejects column projection for binary input.

Memory behavior

The source and destination columns are read in two separate passes.

Each pass:

1. Allocates the final 1-D WholeMemory-backed DistTensor.
2. Uses PyArrow column projection to decode only the selected column.
3. Streams bounded batches directly into each rank’s local partition.

The resulting DistTensor objects are stored directly in the DistMatrix. There is no intermediate (E, 2) tensor, post-load split, unsupported view, or full-size copy.

This trades an additional Parquet scan for substantially lower peak memory usage, which is important for multi-terabyte edge datasets.

Validation

Tests cover:

• Single-column and ordered multi-column projection.
• One-column (E,) shape inference.
• Multiple Parquet files and row groups.
• Reversed physical column order.
• Missing, duplicate, empty, and nonnumeric columns.
• Propagation of device, backend, partition book, expected shape, and dtype policy.
• Direct retention of the two loaded DistTensor allocations by DistMatrix.
• CPU and CUDA destinations.

Local results:

• pylibwholegraph format tests: 15 passed
• structured I/O tests: 21 passed, 8 skipped
• DistMatrix tests: 3 passed, 2 GPU-dependent skipped
• Pre-commit checks: passed

[copy-pr-bot]

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.