gigl.utils.csr#

Memory-lean COO -> CSR/CSC conversion.

graphlearn_torch.utils.coo_to_csr delegates to torch_sparse.SparseTensor, which builds a composite sort key, sorts it, and gathers row and col through the permutation while the caller’s originals are still live. Seven full-size int64 arrays exist at the peak, measured at 7.25x one int64 array on 400M edges. At a multi-billion-edge partition that is hundreds of GiB and the conversion, not the graph, is what exhausts the host.

This implementation is a two-pass counting sort: degrees give the exact output layout up front, so the output is allocated once and written in place, and the input is consumed in chunks so transients are bounded by chunk size rather than by edge count:

peak = row + col + indices + optional edge_ids + 2 x (num_rows + 1)
       + O(chunk + max_degree)

which is 1.5x one int64 array with int32 inputs when edge IDs are omitted, since indices then stays int32, or 2.5x when the int64 edge IDs are retained. The max_degree term comes from within-row sorting; it matters when one row owns a meaningful fraction of the edges.

CompactTopology wraps the result as a GLT Topology.

Attributes#

DEFAULT_BAND_BYTES

DEFAULT_CHUNK_SIZE

DEFAULT_SORT_BLOCK_EDGES

logger

Classes#

CompactTopology

A GLT Topology built by build_csr_from_coo().

CsrBuildResult

CSR arrays and optional COO-position edge IDs in CSR order.

Functions#

build_csr_from_coo(row, col[, num_rows, chunk_size, ...])

Convert a COO edge index to CSR with bounded chunk and sort temporaries.

Module Contents#

class gigl.utils.csr.CompactTopology(edge_index, layout, num_rows=None, *, retain_edge_ids=True)[source]#

Bases: graphlearn_torch.data.Topology

A GLT Topology built by build_csr_from_coo().

Topology.__init__ is not called: it would fabricate a full torch.arange(num_edges) and convert with coo_to_csr. Instead, the counting sort writes original COO positions directly into final CSR order when retained. Edge features need explicit IDs and cannot use this class.

Parameters:

  • edge_index (torch.Tensor) – [2, num_edges] COO. int32 and int64 are used as they are; anything else is cast to int64, as GLT’s Topology does.

  • layout (Literal["CSR", "CSC"]) – The layout GLT samples from.

  • num_rows (Optional[int]) – Size of the compressed dimension. Defaults to the largest compressed id + 1, as GLT’s Topology does.

  • retain_edge_ids (bool) – Retain COO-position IDs for edge sampling by default, at a cost of 8 bytes per edge. Use False only when callers sample without edge IDs.

class gigl.utils.csr.CsrBuildResult[source]#

CSR arrays and optional COO-position edge IDs in CSR order.

edge_ids: torch.Tensor | None = None[source]#
indices: torch.Tensor[source]#
indptr: torch.Tensor[source]#
gigl.utils.csr.build_csr_from_coo(row, col, num_rows=None, chunk_size=DEFAULT_CHUNK_SIZE, sort_block_edges=DEFAULT_SORT_BLOCK_EDGES, band_bytes=DEFAULT_BAND_BYTES, retain_edge_ids=True)[source]#

Convert a COO edge index to CSR with bounded chunk and sort temporaries.

Produces the same CSR arrays as graphlearn_torch.utils.coo_to_csr without its full-size sorting temporaries.

When retain_edge_ids is true, write original COO positions in CSR order alongside columns. The final int64 vector is allocated directly; no full-size input ID vector is needed. Explicit edge IDs and edge weights are unsupported.

To build CSC instead, pass the column indices as row and the row indices as col.

Parameters:

  • row (torch.Tensor) – 1-D row indices, int32 or int64.

  • col (torch.Tensor) – 1-D column indices, same length as row.

  • num_rows (Optional[int]) – Size of the row dimension; indptr has num_rows + 1 entries. Defaults to max(row) + 1, as coo_to_csr does.

  • chunk_size (int) – Edges per scatter chunk. Bounds the transient working set.

  • sort_block_edges (int) – Edges per block of the within-row sort, which orders each row’s columns ascending as coo_to_csr does.

  • band_bytes (int) – Destination window for the scatter, used only when indices could not be placed in memory.

  • retain_edge_ids (bool) – Retain original COO positions in CSR order by default. The int64 IDs cost 8 bytes per edge. Set to False when callers do not sample edge IDs.

Returns:

indptr of shape [num_rows + 1], always int64, and indices of shape [num_edges], int32 when col is int32 and int64 otherwise, plus optional int64 edge_ids of shape [num_edges] in CSR order. Empty input returns an empty ID tensor when retention is enabled and None otherwise.

Return type:

CsrBuildResult

Raises:

ValueError – If row and col disagree in length, a row id is out of range, or the scatter does not exactly fill every row.

gigl.utils.csr.DEFAULT_BAND_BYTES = 4294967296[source]#
gigl.utils.csr.DEFAULT_CHUNK_SIZE = 16777216[source]#
gigl.utils.csr.DEFAULT_SORT_BLOCK_EDGES = 33554432[source]#
gigl.utils.csr.logger[source]#