neighbor_loader.py

# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0

import warnings

from typing import Union, Tuple, Optional, Callable, List, Dict

import numpy as np

import cugraph_pyg
from cugraph_pyg.loader import NodeLoader
from cugraph_pyg.sampler import BaseSampler

from cugraph_pyg.sampler.distributed_sampler import DistributedNeighborSampler
from cugraph_pyg.utils.imports import import_optional

torch_geometric = import_optional("torch_geometric")


class NeighborLoader(NodeLoader):
    """
    Duck-typed version of torch_geometric.loader.NeighborLoader

    Node loader that implements the neighbor sampling
    algorithm used in GraphSAGE.
    """

    def __init__(
        self,
        data: Union[
            "torch_geometric.data.Data",
            "torch_geometric.data.HeteroData",
            Tuple[
                "torch_geometric.data.FeatureStore", "torch_geometric.data.GraphStore"
            ],
        ],
        num_neighbors: Union[
            List[int], Dict["torch_geometric.typing.EdgeType", List[int]]
        ],
        input_nodes: "torch_geometric.typing.InputNodes" = None,
        input_time: "torch_geometric.typing.OptTensor" = None,
        replace: bool = False,
        subgraph_type: Union[
            "torch_geometric.typing.SubgraphType", str
        ] = "directional",
        disjoint: bool = False,
        temporal_strategy: str = "uniform",
        time_attr: Optional[str] = None,
        weight_attr: Optional[str] = None,
        transform: Optional[Callable] = None,
        transform_sampler_output: Optional[Callable] = None,
        is_sorted: bool = False,
        filter_per_worker: Optional[bool] = None,
        neighbor_sampler: Optional["torch_geometric.sampler.NeighborSampler"] = None,
        directed: bool = True,  # Deprecated.
        batch_size: int = 16,
        compression: Optional[str] = None,
        local_seeds_per_call: Optional[int] = None,
        temporal_comparison: Optional[str] = None,
        **kwargs,
    ):
        """
        data: Data, HeteroData, or Tuple[FeatureStore, GraphStore]
            See torch_geometric.loader.NeighborLoader.
        num_neighbors: List[int] or Dict[EdgeType, List[int]]
            Fanout values.
            See torch_geometric.loader.NeighborLoader.
        input_nodes: InputNodes
            Input nodes for sampling.
            See torch_geometric.loader.NeighborLoader.
        input_time: OptTensor (optional)
            See torch_geometric.loader.NeighborLoader.
        replace: bool (optional, default=False)
            Whether to sample with replacement.
            See torch_geometric.loader.NeighborLoader.
        subgraph_type: Union[SubgraphType, str] (optional, default='directional')
            The type of subgraph to return.
            Currently only 'directional' is supported.
            See torch_geometric.loader.NeighborLoader.
        disjoint: bool (optional, default=False)
            Whether to perform disjoint sampling.
            Currently unsupported.
            See torch_geometric.loader.NeighborLoader.
        temporal_strategy: str (optional, default='uniform')
            Currently only 'uniform' is suppported.
            See torch_geometric.loader.NeighborLoader.
        time_attr: str (optional, default=None)
            Used for temporal sampling.
            See torch_geometric.loader.NeighborLoader.
        weight_attr: str (optional, default=None)
            Used for biased sampling.
            See torch_geometric.loader.NeighborLoader.
        transform: Callable (optional, default=None)
            See torch_geometric.loader.NeighborLoader.
        transform_sampler_output: Callable (optional, default=None)
            See torch_geometric.loader.NeighborLoader.
        is_sorted: bool (optional, default=False)
            Ignored by cuGraph.
            See torch_geometric.loader.NeighborLoader.
        filter_per_worker: bool (optional, default=False)
            Currently ignored by cuGraph, but this may
            change once in-memory sampling is implemented.
            See torch_geometric.loader.NeighborLoader.
        neighbor_sampler: torch_geometric.sampler.NeighborSampler
            (optional, default=None)
            Not supported by cuGraph.
            See torch_geometric.loader.NeighborLoader.
        directed: bool (optional, default=True)
            Deprecated.
            See torch_geometric.loader.NeighborLoader.
        batch_size: int (optional, default=16)
            The number of input nodes per output minibatch.
            See torch.utils.dataloader.
        compression: str (optional, default=None)
            The compression type to use if writing samples to disk.
            If not provided, it is automatically chosen.
        local_seeds_per_call: int (optional, default=None)
            The number of seeds to process within a single sampling call.
            Manually tuning this parameter is not recommended but reducing
            it may conserve GPU memory.  The total number of seeds processed
            per sampling call is equal to the sum of this parameter across
            all workers.  If not provided, it will be automatically
            calculated.
            See cugraph_pyg.sampler.BaseDistributedSampler.
        temporal_comparison: str (optional, default='monotonically_decreasing')
            The comparison operator for temporal sampling
            ('strictly_increasing', 'monotonically_increasing',
            'strictly_decreasing', 'monotonically_decreasing', 'last').
            Note that this should be 'last' for temporal_strategy='last'.
            See cugraph_pyg.sampler.BaseDistributedSampler.
        **kwargs
            Other keyword arguments passed to the superclass.
        """

        subgraph_type = torch_geometric.sampler.base.SubgraphType(subgraph_type)

        if temporal_comparison is None:
            temporal_comparison = "monotonically_decreasing"

        if not directed:
            subgraph_type = torch_geometric.sampler.base.SubgraphType.induced
            warnings.warn(
                "The 'directed' argument is deprecated. "
                "Use subgraph_type='induced' instead."
            )
        if subgraph_type != torch_geometric.sampler.base.SubgraphType.directional:
            raise ValueError("Only directional subgraphs are currently supported")
        if disjoint:
            raise ValueError("Disjoint sampling is currently unsupported")
        if temporal_strategy != "uniform":
            warnings.warn("Only the uniform temporal strategy is currently supported")
        if neighbor_sampler is not None:
            raise ValueError("Passing a neighbor sampler is currently unsupported")
        if is_sorted:
            warnings.warn("The 'is_sorted' argument is ignored by cuGraph.")
        if not isinstance(data, (list, tuple)) or not isinstance(
            data[1],
            (cugraph_pyg.data.graph_store.GraphStore,),
        ):
            # Will eventually automatically convert these objects to cuGraph objects.
            raise NotImplementedError("Currently can't accept non-cugraph graphs")

        feature_store, graph_store = data

        if compression is None:
            compression = "CSR" if graph_store.is_homogeneous else "COO"
        elif compression not in ["CSR", "COO"]:
            raise ValueError("Invalid value for compression (expected 'CSR' or 'COO')")

        if not graph_store.is_homogeneous:
            if compression != "COO":
                raise ValueError(
                    "Only COO format is supported for heterogeneous graphs!"
                )

        is_temporal = time_attr is not None

        if is_temporal:
            graph_store._set_time_attr((feature_store, time_attr))

            if input_time is None:
                input_type, input_nodes, _ = (
                    torch_geometric.loader.utils.get_input_nodes(
                        data, input_nodes, None
                    )
                )
                if input_type is None:
                    input_type = list(graph_store._vertex_offsets.keys())[0]
                # will assume the time attribute exists for nodes as well
                input_time = feature_store[input_type, time_attr, None][input_nodes]

        if weight_attr is not None:
            graph_store._set_weight_attr((feature_store, weight_attr))

        if isinstance(num_neighbors, dict):
            sorted_keys, _, _ = graph_store._numeric_edge_types
            fanout_length = len(next(iter(num_neighbors.values())))
            na = np.zeros(fanout_length * len(sorted_keys), dtype="int32")
            for i, key in enumerate(sorted_keys):
                if key in num_neighbors:
                    for hop in range(fanout_length):
                        na[hop * len(sorted_keys) + i] = num_neighbors[key][hop]

            num_neighbors = na

        sampler = BaseSampler(
            DistributedNeighborSampler(
                graph_store._graph,
                retain_original_seeds=True,
                fanout=num_neighbors,
                prior_sources_behavior="exclude",
                deduplicate_sources=True,
                compression=compression,
                compress_per_hop=False,
                with_replacement=replace,
                local_seeds_per_call=local_seeds_per_call,
                biased=(weight_attr is not None),
                heterogeneous=(not graph_store.is_homogeneous),
                temporal=is_temporal,
                temporal_comparison=temporal_comparison,
                vertex_type_offsets=graph_store._vertex_offset_array,
                num_edge_types=len(graph_store.get_all_edge_attrs()),
            ),
            (feature_store, graph_store),
            batch_size=batch_size,
        )

        super().__init__(
            (feature_store, graph_store),
            sampler,
            input_nodes=input_nodes,
            input_time=input_time,
            transform=transform,
            transform_sampler_output=transform_sampler_output,
            filter_per_worker=filter_per_worker,
            batch_size=batch_size,
            **kwargs,
        )