chore: import upstream snapshot with attribution

python/dgl/graphbolt/__init__.py

@@ -0,0 +1,119 @@
+"""Graphbolt."""
+import os
+import sys
+
+from .internal_utils import *
+
+CUDA_ALLOCATOR_ENV_WARNING_STR = """
+An experimental feature for CUDA allocations is turned on for better allocation
+pattern resulting in better memory usage for minibatch GNN training workloads.
+See https://pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf,
+and set the environment variable `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False`
+if you want to disable it and set it True to acknowledge and disable the warning.
+"""
+cuda_allocator_env = os.getenv("PYTORCH_CUDA_ALLOC_CONF")
+WARNING_STR_TO_BE_SHOWN = None
+configs = (
+    {}
+    if cuda_allocator_env is None or len(cuda_allocator_env) == 0
+    else {
+        kv_pair.split(":")[0]: kv_pair.split(":")[1]
+        for kv_pair in cuda_allocator_env.split(",")
+    }
+)
+if "expandable_segments" in configs:
+    if configs["expandable_segments"] != "True":
+        WARNING_STR_TO_BE_SHOWN = (
+            "You should consider `expandable_segments:True` in the"
+            " environment variable `PYTORCH_CUDA_ALLOC_CONF` for lower"
+            " memory usage. See "
+            "https://pytorch.org/docs/stable/notes/cuda.html"
+            "#optimizing-memory-usage-with-pytorch-cuda-alloc-conf"
+        )
+else:
+    configs["expandable_segments"] = "True"
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ",".join(
+        [k + ":" + v for k, v in configs.items()]
+    )
+    WARNING_STR_TO_BE_SHOWN = CUDA_ALLOCATOR_ENV_WARNING_STR
+del configs
+del cuda_allocator_env
+del CUDA_ALLOCATOR_ENV_WARNING_STR
+
+# pylint: disable=wrong-import-position, wrong-import-order
+import torch
+
+### FROM DGL @todo
+from .._ffi import libinfo
+
+
+def load_graphbolt():
+    """Load Graphbolt C++ library"""
+    vers = torch.__version__.split("+", maxsplit=1)[0]
+
+    if sys.platform.startswith("linux"):
+        basename = f"libgraphbolt_pytorch_{vers}.so"
+    elif sys.platform.startswith("darwin"):
+        basename = f"libgraphbolt_pytorch_{vers}.dylib"
+    elif sys.platform.startswith("win"):
+        basename = f"graphbolt_pytorch_{vers}.dll"
+    else:
+        raise NotImplementedError("Unsupported system: %s" % sys.platform)
+
+    dirname = os.path.dirname(libinfo.find_lib_path()[0])
+    path = os.path.join(dirname, "graphbolt", basename)
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"Unable to locate the DGL C++ GraphBolt library at {path}. This "
+            "error typically occurs due to a version mismatch between the "
+            "installed DGL and the PyTorch version you are currently using. "
+            "Please ensure that your DGL installation is compatible with your "
+            "PyTorch version. For more information, refer to the installation "
+            "guide at https://www.dgl.ai/pages/start.html."
+        )
+
+    try:
+        torch.classes.load_library(path)
+    except Exception:  # pylint: disable=W0703
+        raise ImportError("Cannot load Graphbolt C++ library")
+
+
+load_graphbolt()
+
+# pylint: disable=wrong-import-position
+from .base import *
+from .minibatch import *
+from .dataloader import *
+from .datapipes import *
+from .dataset import *
+from .feature_fetcher import *
+from .feature_store import *
+from .impl import *
+from .itemset import *
+from .item_sampler import *
+from .minibatch_transformer import *
+from .negative_sampler import *
+from .sampled_subgraph import *
+from .subgraph_sampler import *
+from .external_utils import add_reverse_edges, exclude_seed_edges
+from .internal import (
+    compact_csc_format,
+    numpy_save_aligned,
+    unique_and_compact,
+    unique_and_compact_csc_formats,
+)
+
+if torch.cuda.is_available() and not built_with_cuda():
+    raise ImportError(
+        "torch was installed with CUDA support while GraphBolt's CPU version "
+        "is installed. Consider reinstalling GraphBolt with CUDA support, see "
+        "installation instructions at https://www.dgl.ai/pages/start.html"
+    )
+
+if torch.cuda.is_available() and WARNING_STR_TO_BE_SHOWN is not None:
+    gb_warning(WARNING_STR_TO_BE_SHOWN)
+del WARNING_STR_TO_BE_SHOWN
+
+torch.ops.graphbolt.set_num_io_uring_threads(
+    min((torch.get_num_threads() + 1) // 2, 8)
+)

python/dgl/graphbolt/base.py

@@ -0,0 +1,524 @@
+"""Base types and utilities for Graph Bolt."""
+
+from collections import deque
+from dataclasses import dataclass
+
+import torch
+from torch.torch_version import TorchVersion
+
+if (
+    TorchVersion(torch.__version__) >= "2.3.0"
+    and TorchVersion(torch.__version__) < "2.3.1"
+):
+    # Due to https://github.com/dmlc/dgl/issues/7380, for torch 2.3.0, we need
+    # to check if dill is available before using it.
+    torch.utils.data.datapipes.utils.common.DILL_AVAILABLE = (
+        torch.utils._import_utils.dill_available()
+    )
+
+# pylint: disable=wrong-import-position
+from torch.utils.data import functional_datapipe, IterDataPipe
+
+from .internal_utils import (
+    get_nonproperty_attributes,
+    recursive_apply,
+    recursive_apply_reduce_all,
+)
+
+__all__ = [
+    "CANONICAL_ETYPE_DELIMITER",
+    "ORIGINAL_EDGE_ID",
+    "etype_str_to_tuple",
+    "etype_tuple_to_str",
+    "CopyTo",
+    "Waiter",
+    "Bufferer",
+    "EndMarker",
+    "isin",
+    "index_select",
+    "expand_indptr",
+    "indptr_edge_ids",
+    "CSCFormatBase",
+    "seed",
+    "seed_type_str_to_ntypes",
+    "get_host_to_device_uva_stream",
+    "get_device_to_host_uva_stream",
+]
+
+CANONICAL_ETYPE_DELIMITER = ":"
+ORIGINAL_EDGE_ID = "_ORIGINAL_EDGE_ID"
+
+
+# There needs to be a single instance of the uva_stream, if it is created
+# multiple times, it leads to multiple CUDA memory pools and memory leaks.
+def get_host_to_device_uva_stream():
+    """The host to device copy stream to be used for pipeline parallelism."""
+    if not hasattr(get_host_to_device_uva_stream, "stream"):
+        get_host_to_device_uva_stream.stream = torch.cuda.Stream(priority=-1)
+    return get_host_to_device_uva_stream.stream
+
+
+def get_device_to_host_uva_stream():
+    """The device to host copy stream to be used for pipeline parallelism."""
+    if not hasattr(get_device_to_host_uva_stream, "stream"):
+        get_device_to_host_uva_stream.stream = torch.cuda.Stream(priority=-1)
+    return get_device_to_host_uva_stream.stream
+
+
+def seed(val):
+    """Set the random seed of Graphbolt.
+
+    Parameters
+    ----------
+    val : int
+        The seed.
+    """
+    torch.ops.graphbolt.set_seed(val)
+
+
+def isin(elements, test_elements):
+    """Tests if each element of elements is in test_elements. Returns a boolean
+    tensor of the same shape as elements that is True for elements in
+    test_elements and False otherwise.
+
+    Parameters
+    ----------
+    elements : torch.Tensor
+        A 1D tensor represents the input elements.
+    test_elements : torch.Tensor
+        A 1D tensor represents the values to test against for each input.
+
+    Examples
+    --------
+    >>> isin(torch.tensor([1, 2, 3, 4]), torch.tensor([2, 3]))
+    tensor([[False,  True,  True,  False]])
+    """
+    assert elements.dim() == 1, "Elements should be 1D tensor."
+    assert test_elements.dim() == 1, "Test_elements should be 1D tensor."
+    return torch.ops.graphbolt.isin(elements, test_elements)
+
+
+if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
+
+    torch_fake_decorator = (
+        torch.library.impl_abstract
+        if TorchVersion(torch.__version__) < TorchVersion("2.4.0a0")
+        else torch.library.register_fake
+    )
+
+    @torch_fake_decorator("graphbolt::expand_indptr")
+    def expand_indptr_fake(indptr, dtype, node_ids, output_size):
+        """Fake implementation of expand_indptr for torch.compile() support."""
+        if output_size is None:
+            output_size = torch.library.get_ctx().new_dynamic_size()
+        if dtype is None:
+            dtype = node_ids.dtype
+        return indptr.new_empty(output_size, dtype=dtype)
+
+
+def expand_indptr(indptr, dtype=None, node_ids=None, output_size=None):
+    """Converts a given indptr offset tensor to a COO format tensor. If
+    node_ids is not given, it is assumed to be equal to
+    torch.arange(indptr.size(0) - 1, dtype=dtype, device=indptr.device).
+
+    This is equivalent to
+
+    .. code:: python
+
+       if node_ids is None:
+           node_ids = torch.arange(len(indptr) - 1, dtype=dtype, device=indptr.device)
+       return node_ids.to(dtype).repeat_interleave(indptr.diff())
+
+    Parameters
+    ----------
+    indptr : torch.Tensor
+        A 1D tensor represents the csc_indptr tensor.
+    dtype : Optional[torch.dtype]
+        The dtype of the returned output tensor.
+    node_ids : Optional[torch.Tensor]
+        A 1D tensor represents the column node ids that the returned tensor will
+        be populated with.
+    output_size : Optional[int]
+        The size of the output tensor. Should be equal to indptr[-1]. Using this
+        argument avoids a stream synchronization to calculate the output shape.
+
+    Returns
+    -------
+    torch.Tensor
+        The converted COO tensor with values from node_ids.
+    """
+    assert indptr.dim() == 1, "Indptr should be 1D tensor."
+    assert not (
+        node_ids is None and dtype is None
+    ), "One of node_ids or dtype must be given."
+    assert (
+        node_ids is None or node_ids.dim() == 1
+    ), "Node_ids should be 1D tensor."
+    if dtype is None:
+        dtype = node_ids.dtype
+    return torch.ops.graphbolt.expand_indptr(
+        indptr, dtype, node_ids, output_size
+    )
+
+
+if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
+
+    torch_fake_decorator = (
+        torch.library.impl_abstract
+        if TorchVersion(torch.__version__) < TorchVersion("2.4.0a0")
+        else torch.library.register_fake
+    )
+
+    @torch_fake_decorator("graphbolt::indptr_edge_ids")
+    def indptr_edge_ids_fake(indptr, dtype, offset, output_size):
+        """Fake implementation of indptr_edge_ids for torch.compile() support."""
+        if output_size is None:
+            output_size = torch.library.get_ctx().new_dynamic_size()
+        if dtype is None:
+            dtype = offset.dtype
+        return indptr.new_empty(output_size, dtype=dtype)
+
+
+def indptr_edge_ids(indptr, dtype=None, offset=None, output_size=None):
+    """Converts a given indptr offset tensor to a COO format tensor for the edge
+    ids. For a given indptr [0, 2, 5, 7] and offset tensor [0, 100, 200], the
+    output will be [0, 1, 100, 101, 102, 201, 202]. If offset was not provided,
+    the output would be [0, 1, 0, 1, 2, 0, 1].
+
+    Parameters
+    ----------
+    indptr : torch.Tensor
+        A 1D tensor represents the csc_indptr tensor.
+    dtype : Optional[torch.dtype]
+        The dtype of the returned output tensor.
+    offset : Optional[torch.Tensor]
+        A 1D tensor represents the offsets that the returned tensor will be
+        populated with.
+    output_size : Optional[int]
+        The size of the output tensor. Should be equal to indptr[-1]. Using this
+        argument avoids a stream synchronization to calculate the output shape.
+
+    Returns
+    -------
+    torch.Tensor
+        The converted COO edge ids tensor.
+    """
+    assert indptr.dim() == 1, "Indptr should be 1D tensor."
+    assert offset is None or offset.dim() == 1, "Offset should be 1D tensor."
+    if dtype is None:
+        dtype = offset.dtype
+    return torch.ops.graphbolt.indptr_edge_ids(
+        indptr, dtype, offset, output_size
+    )
+
+
+def index_select(tensor, index):
+    """Returns a new tensor which indexes the input tensor along dimension dim
+    using the entries in index.
+
+    The returned tensor has the same number of dimensions as the original tensor
+    (tensor). The first dimension has the same size as the length of index;
+    other dimensions have the same size as in the original tensor.
+
+    When tensor is a pinned tensor and index.is_cuda is True, the operation runs
+    on the CUDA device and the returned tensor will also be on CUDA.
+
+    Parameters
+    ----------
+    tensor : torch.Tensor
+        The input tensor.
+    index : torch.Tensor
+        The 1-D tensor containing the indices to index.
+
+    Returns
+    -------
+    torch.Tensor
+        The indexed input tensor, equivalent to tensor[index]. If index is in
+        pinned memory, then the result is placed into pinned memory as well.
+    """
+    assert index.dim() == 1, "Index should be 1D tensor."
+    return torch.ops.graphbolt.index_select(tensor, index)
+
+
+def etype_tuple_to_str(c_etype):
+    """Convert canonical etype from tuple to string.
+
+    Examples
+    --------
+    >>> c_etype = ("user", "like", "item")
+    >>> c_etype_str = _etype_tuple_to_str(c_etype)
+    >>> print(c_etype_str)
+    "user:like:item"
+    """
+    assert isinstance(c_etype, tuple) and len(c_etype) == 3, (
+        "Passed-in canonical etype should be in format of (str, str, str). "
+        f"But got {c_etype}."
+    )
+    return CANONICAL_ETYPE_DELIMITER.join(c_etype)
+
+
+def etype_str_to_tuple(c_etype):
+    """Convert canonical etype from string to tuple.
+
+    Examples
+    --------
+    >>> c_etype_str = "user:like:item"
+    >>> c_etype = _etype_str_to_tuple(c_etype_str)
+    >>> print(c_etype)
+    ("user", "like", "item")
+    """
+    if isinstance(c_etype, tuple):
+        return c_etype
+    ret = tuple(c_etype.split(CANONICAL_ETYPE_DELIMITER))
+    assert len(ret) == 3, (
+        "Passed-in canonical etype should be in format of 'str:str:str'. "
+        f"But got {c_etype}."
+    )
+    return ret
+
+
+def seed_type_str_to_ntypes(seed_type, seed_size):
+    """Convert seeds type to node types from string to list.
+
+    Examples
+    --------
+    1. node pairs
+
+    >>> seed_type = "user:like:item"
+    >>> seed_size = 2
+    >>> node_type = seed_type_str_to_ntypes(seed_type, seed_size)
+    >>> print(node_type)
+    ["user", "item"]
+
+    2. hyperlink
+
+    >>> seed_type = "query:user:item"
+    >>> seed_size = 3
+    >>> node_type = seed_type_str_to_ntypes(seed_type, seed_size)
+    >>> print(node_type)
+    ["query", "user", "item"]
+    """
+    assert isinstance(
+        seed_type, str
+    ), f"Passed-in seed type should be string, but got {type(seed_type)}"
+    ntypes = seed_type.split(CANONICAL_ETYPE_DELIMITER)
+    is_hyperlink = len(ntypes) == seed_size
+    if not is_hyperlink:
+        ntypes = ntypes[::2]
+    return ntypes
+
+
+def apply_to(x, device, non_blocking=False):
+    """Apply `to` function to object x only if it has `to`."""
+
+    if device == "pinned" and hasattr(x, "pin_memory"):
+        return x.pin_memory()
+    if not hasattr(x, "to"):
+        return x
+    if not non_blocking:
+        return x.to(device)
+    return x.to(device, non_blocking=True)
+
+
+def is_object_pinned(obj):
+    """Recursively check all members of the object and return True if only if
+    all are pinned."""
+
+    for attr in get_nonproperty_attributes(obj):
+        member_result = recursive_apply_reduce_all(
+            getattr(obj, attr),
+            lambda x: x is None or x.is_pinned(),
+        )
+        if not member_result:
+            return False
+    return True
+
+
+@functional_datapipe("copy_to")
+class CopyTo(IterDataPipe):
+    """DataPipe that transfers each element yielded from the previous DataPipe
+    to the given device. For MiniBatch, only the related attributes
+    (automatically inferred) will be transferred by default.
+
+    Functional name: :obj:`copy_to`.
+
+    When ``data`` has ``to`` method implemented, ``CopyTo`` will be equivalent
+    to
+
+    .. code:: python
+
+       for data in datapipe:
+           yield data.to(device)
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The DataPipe.
+    device : torch.device
+        The PyTorch CUDA device.
+    non_blocking : bool
+        Whether the copy should be performed without blocking. All elements have
+        to be already in pinned system memory if enabled. Default is False.
+    """
+
+    def __init__(self, datapipe, device, non_blocking=False):
+        super().__init__()
+        self.datapipe = datapipe
+        self.device = torch.device(device)
+        self.non_blocking = non_blocking
+
+    def __iter__(self):
+        for data in self.datapipe:
+            yield recursive_apply(
+                data, apply_to, self.device, self.non_blocking
+            )
+
+
+@functional_datapipe("mark_end")
+class EndMarker(IterDataPipe):
+    """Used to mark the end of a datapipe and is a no-op."""
+
+    def __init__(self, datapipe):
+        self.datapipe = datapipe
+
+    def __iter__(self):
+        yield from self.datapipe
+
+
+@functional_datapipe("buffer")
+class Bufferer(IterDataPipe):
+    """Buffers items before yielding them.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The data pipeline.
+    buffer_size : int, optional
+        The size of the buffer which stores the fetched samples. If data coming
+        from datapipe has latency spikes, consider setting to a higher value.
+        Default is 1.
+    """
+
+    def __init__(self, datapipe, buffer_size=1):
+        self.datapipe = datapipe
+        if buffer_size <= 0:
+            raise ValueError(
+                "'buffer_size' is required to be a positive integer."
+            )
+        self.buffer = deque(maxlen=buffer_size)
+
+    def __iter__(self):
+        for data in self.datapipe:
+            if len(self.buffer) < self.buffer.maxlen:
+                self.buffer.append(data)
+            else:
+                return_data = self.buffer.popleft()
+                self.buffer.append(data)
+                yield return_data
+        while len(self.buffer) > 0:
+            yield self.buffer.popleft()
+
+    def __getstate__(self):
+        state = (self.datapipe, self.buffer.maxlen)
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        self.datapipe, buffer_size = state
+        self.buffer = deque(maxlen=buffer_size)
+
+    def reset(self):
+        """Resets the state of the datapipe."""
+        self.buffer.clear()
+
+
+@functional_datapipe("wait")
+class Waiter(IterDataPipe):
+    """Calls the wait function of all items."""
+
+    def __init__(self, datapipe):
+        self.datapipe = datapipe
+
+    def __iter__(self):
+        for data in self.datapipe:
+            data.wait()
+            yield data
+
+
+@dataclass
+class CSCFormatBase:
+    r"""Basic class representing data in Compressed Sparse Column (CSC) format.
+
+    Examples
+    --------
+    >>> indptr = torch.tensor([0, 1, 3])
+    >>> indices = torch.tensor([1, 4, 2])
+    >>> csc_foramt_base = CSCFormatBase(indptr=indptr, indices=indices)
+    >>> print(csc_format_base.indptr)
+    ... torch.tensor([0, 1, 3])
+    >>> print(csc_foramt_base)
+    ... torch.tensor([1, 4, 2])
+    """
+
+    indptr: torch.Tensor = None
+    indices: torch.Tensor = None
+
+    def __init__(self, indptr: torch.Tensor, indices: torch.Tensor):
+        self.indptr = indptr
+        self.indices = indices
+        if not indptr.is_cuda:
+            assert self.indptr[-1] == len(
+                self.indices
+            ), "The last element of indptr should be the same as the length of indices."
+
+    def __repr__(self) -> str:
+        return _csc_format_base_str(self)
+
+    def to(  # pylint: disable=invalid-name
+        self, device: torch.device, non_blocking=False
+    ) -> None:
+        """Copy `CSCFormatBase` to the specified device using reflection."""
+
+        for attr in dir(self):
+            # Only copy member variables.
+            if not callable(getattr(self, attr)) and not attr.startswith("__"):
+                setattr(
+                    self,
+                    attr,
+                    recursive_apply(
+                        getattr(self, attr),
+                        apply_to,
+                        device,
+                        non_blocking=non_blocking,
+                    ),
+                )
+
+        return self
+
+    def pin_memory(self):
+        """Copy `SampledSubgraph` to the pinned memory using reflection."""
+
+        return self.to("pinned")
+
+    def is_pinned(self) -> bool:
+        """Check whether `SampledSubgraph` is pinned using reflection."""
+
+        return is_object_pinned(self)
+
+
+def _csc_format_base_str(csc_format_base: CSCFormatBase) -> str:
+    final_str = "CSCFormatBase("
+
+    def _add_indent(_str, indent):
+        lines = _str.split("\n")
+        lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
+        return "\n".join(lines)
+
+    final_str += (
+        f"indptr={_add_indent(str(csc_format_base.indptr), 21)},\n" + " " * 14
+    )
+    final_str += (
+        f"indices={_add_indent(str(csc_format_base.indices), 22)},\n" + ")"
+    )
+    return final_str

python/dgl/graphbolt/dataloader.py

@@ -0,0 +1,191 @@
+"""Graph Bolt DataLoaders"""
+
+import torch
+import torch.utils.data as torch_data
+
+from .base import CopyTo
+from .datapipes import (
+    datapipe_graph_to_adjlist,
+    find_dps,
+    replace_dp,
+    traverse_dps,
+)
+from .feature_fetcher import FeatureFetcher, FeatureFetcherStartMarker
+from .impl.neighbor_sampler import SamplePerLayer
+from .internal_utils import gb_warning
+from .item_sampler import ItemSampler
+from .minibatch_transformer import MiniBatchTransformer
+
+
+__all__ = [
+    "DataLoader",
+]
+
+
+def _find_and_wrap_parent(datapipe_graph, target_datapipe, wrapper, **kwargs):
+    """Find parent of target_datapipe and wrap it with ."""
+    datapipes = find_dps(
+        datapipe_graph,
+        target_datapipe,
+    )
+    datapipe_adjlist = datapipe_graph_to_adjlist(datapipe_graph)
+    for datapipe in datapipes:
+        datapipe_id = id(datapipe)
+        for parent_datapipe_id in datapipe_adjlist[datapipe_id][1]:
+            parent_datapipe, _ = datapipe_adjlist[parent_datapipe_id]
+            datapipe_graph = replace_dp(
+                datapipe_graph,
+                parent_datapipe,
+                wrapper(parent_datapipe, **kwargs),
+            )
+    return datapipe_graph
+
+
+def _set_worker_id(worked_id):
+    torch.ops.graphbolt.set_worker_id(worked_id)
+
+
+class MultiprocessingWrapper(torch_data.IterDataPipe):
+    """Wraps a datapipe with multiprocessing.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The data pipeline.
+    num_workers : int, optional
+        The number of worker processes. Default is 0, meaning that there
+        will be no multiprocessing.
+    persistent_workers : bool, optional
+        If True, the data loader will not shut down the worker processes after a
+        dataset has been consumed once. This allows to maintain the workers
+        instances alive.
+    """
+
+    def __init__(self, datapipe, num_workers=0, persistent_workers=True):
+        self.datapipe = datapipe
+        self.dataloader = torch_data.DataLoader(
+            datapipe,
+            batch_size=None,
+            num_workers=num_workers,
+            persistent_workers=(num_workers > 0) and persistent_workers,
+            worker_init_fn=_set_worker_id if num_workers > 0 else None,
+        )
+
+    def __iter__(self):
+        yield from self.dataloader
+
+
+class DataLoader(MiniBatchTransformer):
+    """Multiprocessing DataLoader.
+
+    Iterates over the data pipeline with everything before feature fetching
+    (i.e. :class:`dgl.graphbolt.FeatureFetcher`) in subprocesses, and
+    everything after feature fetching in the main process. The datapipe
+    is modified in-place as a result.
+
+    When the copy_to operation is placed earlier in the data pipeline, the
+    num_workers argument is required to be 0 as utilizing CUDA in multiple
+    worker processes is not supported.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The data pipeline.
+    num_workers : int, optional
+        Number of worker processes. Default is 0.
+    persistent_workers : bool, optional
+        If True, the data loader will not shut down the worker processes after a
+        dataset has been consumed once. This allows to maintain the workers
+        instances alive.
+    max_uva_threads : int, optional
+        Limits the number of CUDA threads used for UVA copies so that the rest
+        of the computations can run simultaneously with it. Setting it to a too
+        high value will limit the amount of overlap while setting it too low may
+        cause the PCI-e bandwidth to not get fully utilized. Manually tuned
+        default is 10240, meaning around 5-7 Streaming Multiprocessors.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        num_workers=0,
+        persistent_workers=True,
+        max_uva_threads=10240,
+    ):
+        # Multiprocessing requires two modifications to the datapipe:
+        #
+        # 1. Insert a stage after ItemSampler to distribute the
+        #    minibatches evenly across processes.
+        # 2. Cut the datapipe at FeatureFetcher, and wrap the inner datapipe
+        #    of the FeatureFetcher with a multiprocessing PyTorch DataLoader.
+
+        datapipe = datapipe.mark_end()
+        datapipe_graph = traverse_dps(datapipe)
+
+        if num_workers > 0:
+            # (1) Insert minibatch distribution.
+            # TODO(BarclayII): Currently I'm using sharding_filter() as a
+            # concept demonstration. Later on minibatch distribution should be
+            # merged into ItemSampler to maximize efficiency.
+            item_samplers = find_dps(
+                datapipe_graph,
+                ItemSampler,
+            )
+            for item_sampler in item_samplers:
+                datapipe_graph = replace_dp(
+                    datapipe_graph,
+                    item_sampler,
+                    item_sampler.sharding_filter(),
+                )
+
+            # (2) Cut datapipe at FeatureFetcher and wrap.
+            datapipe_graph = _find_and_wrap_parent(
+                datapipe_graph,
+                FeatureFetcherStartMarker,
+                MultiprocessingWrapper,
+                num_workers=num_workers,
+                persistent_workers=persistent_workers,
+            )
+
+        # (3) Limit the number of UVA threads used if the feature_fetcher
+        # or any of the samplers have overlapping optimization enabled.
+        if num_workers == 0 and torch.cuda.is_available():
+            feature_fetchers = find_dps(
+                datapipe_graph,
+                FeatureFetcher,
+            )
+            for feature_fetcher in feature_fetchers:
+                if feature_fetcher.max_num_stages > 0:  # Overlap enabled.
+                    torch.ops.graphbolt.set_max_uva_threads(max_uva_threads)
+
+        if num_workers == 0 and torch.cuda.is_available():
+            samplers = find_dps(
+                datapipe_graph,
+                SamplePerLayer,
+            )
+            for sampler in samplers:
+                if sampler.overlap_fetch:
+                    torch.ops.graphbolt.set_max_uva_threads(max_uva_threads)
+
+        # (4) Cut datapipe at CopyTo and wrap with pinning and prefetching
+        # before it. This enables enables non_blocking copies to the device.
+        # Prefetching enables the data pipeline up to the CopyTo to run in a
+        # separate thread.
+        copiers = find_dps(datapipe_graph, CopyTo)
+        if len(copiers) > 1:
+            gb_warning(
+                "Multiple CopyTo operations were found in the datapipe graph."
+                " This case is not officially supported."
+            )
+        for copier in copiers:
+            # We enable the prefetch at all times for good CPU only performance.
+            datapipe_graph = replace_dp(
+                datapipe_graph,
+                copier,
+                # Add prefetch so that CPU and GPU can run concurrently.
+                copier.datapipe.prefetch(2).copy_to(
+                    copier.device, non_blocking=True
+                ),
+            )
+
+        super().__init__(datapipe)

python/dgl/graphbolt/datapipes/__init__.py

@@ -0,0 +1,3 @@
+"""GraphBolt's datapipes, mostly copied from "torchdata==0.7.1"."""
+from .utils import *
+from .visualization import *

python/dgl/graphbolt/datapipes/utils.py

@@ -0,0 +1,372 @@
+"""DataPipe utilities"""
+
+import threading
+import time
+
+from collections import deque
+from typing import final, List, Set, Type  # pylint: disable=no-name-in-module
+
+from torch.utils.data import functional_datapipe, IterDataPipe, MapDataPipe
+from torch.utils.data.graph import DataPipe, DataPipeGraph, traverse_dps
+
+__all__ = [
+    "datapipe_graph_to_adjlist",
+    "find_dps",
+    "replace_dp",
+    "traverse_dps",
+]
+
+# Copied from:
+# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/datapipes/iter/util/prefetcher.py#L19-L20
+# Interval between buffer fulfillment checks
+PRODUCER_SLEEP_INTERVAL = 0.0001
+# Interval between checking items availability in buffer
+CONSUMER_SLEEP_INTERVAL = 0.0001
+
+
+def _get_parents(result_dict, datapipe_graph):
+    for k, (v, parents) in datapipe_graph.items():
+        if k not in result_dict:
+            result_dict[k] = (v, list(parents.keys()))
+            _get_parents(result_dict, parents)
+
+
+def datapipe_graph_to_adjlist(datapipe_graph):
+    """Given a DataPipe graph returned by
+    :func:`torch.utils.data.graph.traverse_dps` in DAG form, convert it into
+    adjacency list form.
+
+    Namely, :func:`torch.utils.data.graph.traverse_dps` returns the following
+    data structure:
+
+    .. code::
+
+       {
+           id(datapipe): (
+               datapipe,
+               {
+                   id(parent1_of_datapipe): (parent1_of_datapipe, {...}),
+                   id(parent2_of_datapipe): (parent2_of_datapipe, {...}),
+                   ...
+               }
+           )
+       }
+
+    We convert it into the following for easier access:
+
+    .. code::
+
+       {
+           id(datapipe1): (
+               datapipe1,
+               [id(parent1_of_datapipe1), id(parent2_of_datapipe1), ...]
+           ),
+           id(datapipe2): (
+               datapipe2,
+               [id(parent1_of_datapipe2), id(parent2_of_datapipe2), ...]
+           ),
+           ...
+       }
+    """
+
+    result_dict = {}
+    _get_parents(result_dict, datapipe_graph)
+    return result_dict
+
+
+# Copied from:
+# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/dataloader2/graph/utils.py#L16-L35
+def find_dps(graph: DataPipeGraph, dp_type: Type[DataPipe]) -> List[DataPipe]:
+    r"""
+    Given the graph of DataPipe generated by ``traverse_dps`` function, return DataPipe
+    instances with the provided DataPipe type.
+    """
+    dps: List[DataPipe] = []
+    cache: Set[int] = set()
+
+    def helper(g) -> None:  # pyre-ignore
+        for dp_id, (dp, src_graph) in g.items():
+            if dp_id in cache:
+                continue
+            cache.add(dp_id)
+            # Please not use `isinstance`, there is a bug.
+            if type(dp) is dp_type:  # pylint: disable=unidiomatic-typecheck
+                dps.append(dp)
+            helper(src_graph)
+
+    helper(graph)
+
+    return dps
+
+
+# Copied from:
+# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/dataloader2/graph/utils.py#L82-L97
+# Given the DataPipe needs to be replaced and the expected DataPipe, return a new graph
+def replace_dp(
+    graph: DataPipeGraph, old_datapipe: DataPipe, new_datapipe: DataPipe
+) -> DataPipeGraph:
+    r"""
+    Given the graph of DataPipe generated by ``traverse_dps`` function and the
+    DataPipe to be replaced and the new DataPipe, return the new graph of
+    DataPipe.
+    """
+    assert len(graph) == 1
+
+    if id(old_datapipe) in graph:
+        graph = traverse_dps(new_datapipe)
+
+    final_datapipe = list(graph.values())[0][0]
+
+    for recv_dp, send_graph in graph.values():
+        _replace_dp(recv_dp, send_graph, old_datapipe, new_datapipe)
+
+    return traverse_dps(final_datapipe)
+
+
+# For each `recv_dp`, find if the source_datapipe needs to be replaced by the new one.
+# If found, find where the `old_dp` is located in `recv_dp` and switch it to the `new_dp`
+def _replace_dp(
+    recv_dp, send_graph: DataPipeGraph, old_dp: DataPipe, new_dp: DataPipe
+) -> None:
+    old_dp_id = id(old_dp)
+    for send_id in send_graph:
+        if send_id == old_dp_id:
+            _assign_attr(recv_dp, old_dp, new_dp, inner_dp=True)
+        else:
+            send_dp, sub_send_graph = send_graph[send_id]
+            _replace_dp(send_dp, sub_send_graph, old_dp, new_dp)
+
+
+# Recursively re-assign datapipe for the sake of nested data structure
+# `inner_dp` is used to prevent recursive call if we have already met a `DataPipe`
+def _assign_attr(obj, old_dp, new_dp, inner_dp: bool = False):
+    if obj is old_dp:
+        return new_dp
+    elif isinstance(obj, (IterDataPipe, MapDataPipe)):
+        # Prevent recursive call for DataPipe
+        if not inner_dp:
+            return None
+        for k in list(obj.__dict__.keys()):
+            new_obj = _assign_attr(obj.__dict__[k], old_dp, new_dp)
+            if new_obj is not None:
+                obj.__dict__[k] = new_obj
+                break
+        return None
+    elif isinstance(obj, dict):
+        for k in list(obj.keys()):
+            new_obj = _assign_attr(obj[k], old_dp, new_dp)
+            if new_obj is not None:
+                obj[k] = new_obj
+                break
+        return None
+    # Tuple is immutable, has to re-create a tuple
+    elif isinstance(obj, tuple):
+        temp_list = []
+        flag = False
+        for item in obj:
+            new_obj = _assign_attr(item, old_dp, new_dp, inner_dp)
+            if new_obj is not None:
+                flag = True
+                temp_list.append(new_dp)
+            else:
+                temp_list.append(item)
+        if flag:
+            return tuple(temp_list)  # Special case
+        else:
+            return None
+    elif isinstance(obj, list):
+        for i in range(len(obj)):  # pylint: disable=consider-using-enumerate
+            new_obj = _assign_attr(obj[i], old_dp, new_dp, inner_dp)
+            if new_obj is not None:
+                obj[i] = new_obj
+                break
+        return None
+    elif isinstance(obj, set):
+        new_obj = None
+        for item in obj:
+            if _assign_attr(item, old_dp, new_dp, inner_dp) is not None:
+                new_obj = new_dp
+                break
+        if new_obj is not None:
+            obj.remove(old_dp)
+            obj.add(new_dp)
+        return None
+    else:
+        return None
+
+
+class _PrefetchData:
+    def __init__(self, source_datapipe, buffer_size: int):
+        self.run_prefetcher: bool = True
+        self.prefetch_buffer: Deque = deque()
+        self.buffer_size: int = buffer_size
+        self.source_datapipe = source_datapipe
+        self.stop_iteration: bool = False
+        self.paused: bool = False
+
+
+# Copied from:
+# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/datapipes/iter/util/prefetcher.py#L34-L172
+@functional_datapipe("prefetch")
+class PrefetcherIterDataPipe(IterDataPipe):
+    r"""
+    Prefetches elements from the source DataPipe and puts them into a buffer
+    (functional name: ``prefetch``). Prefetching performs the operations (e.g.
+    I/O, computations) of the DataPipes up to this one ahead of time and stores
+    the result in the buffer, ready to be consumed by the subsequent DataPipe.
+    It has no effect aside from getting the sample ready ahead of time.
+
+    This is used by ``MultiProcessingReadingService`` when the arguments
+    ``worker_prefetch_cnt`` (for prefetching at each worker process) or
+    ``main_prefetch_cnt`` (for prefetching at the main loop) are greater than 0.
+
+    Beyond the built-in use cases, this can be useful to put after I/O DataPipes
+    that have expensive I/O operations (e.g. takes a long time to request a file
+    from a remote server).
+
+    Args:
+        source_datapipe: IterDataPipe from which samples are prefetched
+        buffer_size: the size of the buffer which stores the prefetched samples
+
+    Example:
+        >>> from torchdata.datapipes.iter import IterableWrapper
+        >>> dp = IterableWrapper(file_paths).open_files().prefetch(5)
+    """
+
+    def __init__(self, source_datapipe, buffer_size: int = 10):
+        self.source_datapipe = source_datapipe
+        if buffer_size <= 0:
+            raise ValueError(
+                "'buffer_size' is required to be a positive integer."
+            )
+        self.buffer_size = buffer_size
+        self.thread: Optional[threading.Thread] = None
+        self.prefetch_data: Optional[_PrefetchData] = None
+
+    @staticmethod
+    def thread_worker(
+        prefetch_data: _PrefetchData,
+    ):  # pylint: disable=missing-function-docstring
+        itr = iter(prefetch_data.source_datapipe)
+        while not prefetch_data.stop_iteration:
+            # Run if not paused
+            while prefetch_data.run_prefetcher:
+                if (
+                    len(prefetch_data.prefetch_buffer)
+                    < prefetch_data.buffer_size
+                ):
+                    try:
+                        item = next(itr)
+                        prefetch_data.prefetch_buffer.append(item)
+                    except Exception as e:  # pylint: disable=broad-except
+                        prefetch_data.run_prefetcher = False
+                        prefetch_data.stop_iteration = True
+                        prefetch_data.prefetch_buffer.append(e)
+                else:  # Buffer is full, waiting for main thread to consume items
+                    # TODO: Calculate sleep interval based on previous consumption speed
+                    time.sleep(PRODUCER_SLEEP_INTERVAL)
+            prefetch_data.paused = True
+            # Sleep longer when this prefetcher thread is paused
+            time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
+
+    def __iter__(self):
+        try:
+            prefetch_data = _PrefetchData(
+                self.source_datapipe, self.buffer_size
+            )
+            self.prefetch_data = prefetch_data
+            thread = threading.Thread(
+                target=PrefetcherIterDataPipe.thread_worker,
+                args=(prefetch_data,),
+                daemon=True,
+            )
+            thread.start()
+            self.thread = thread
+
+            while (
+                not prefetch_data.stop_iteration
+                or len(prefetch_data.prefetch_buffer) > 0
+            ):
+                if len(prefetch_data.prefetch_buffer) > 0:
+                    data = prefetch_data.prefetch_buffer.popleft()
+                    if isinstance(data, Exception):
+                        if isinstance(data, StopIteration):
+                            break
+                        raise data
+                    yield data
+                else:
+                    time.sleep(CONSUMER_SLEEP_INTERVAL)
+        finally:
+            if "prefetch_data" in locals():
+                prefetch_data.run_prefetcher = False
+                prefetch_data.stop_iteration = True
+                prefetch_data.paused = False
+            if "thread" in locals():
+                thread.join()
+
+    def __getstate__(self):
+        """
+        Getting state in threading environment requires next operations:
+            1) Stopping of the producer thread.
+            2) Saving buffer.
+            3) Adding lazy restart of producer thread when __next__ is called again
+              (this will guarantee that you only change state of the source_datapipe
+               after entire state of the graph is saved).
+        """
+        # TODO: Update __getstate__ and __setstate__ to support snapshotting and restoration
+        return {
+            "source_datapipe": self.source_datapipe,
+            "buffer_size": self.buffer_size,
+        }
+
+    def __setstate__(self, state):
+        self.source_datapipe = state["source_datapipe"]
+        self.buffer_size = state["buffer_size"]
+        self.thread = None
+
+    @final
+    def reset(self):  # pylint: disable=missing-function-docstring
+        self.shutdown()
+
+    def pause(self):  # pylint: disable=missing-function-docstring
+        if self.thread is not None:
+            assert self.prefetch_data is not None
+            self.prefetch_data.run_prefetcher = False
+            if self.thread.is_alive():
+                # Blocking until the thread is paused
+                while not self.prefetch_data.paused:
+                    time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
+
+    @final
+    def resume(self):  # pylint: disable=missing-function-docstring
+        if (
+            self.thread is not None
+            and self.prefetch_data is not None
+            and (
+                not self.prefetch_data.stop_iteration
+                or len(self.prefetch_data.prefetch_buffer) > 0
+            )
+        ):
+            self.prefetch_data.run_prefetcher = True
+            self.prefetch_data.paused = False
+
+    @final
+    def shutdown(self):  # pylint: disable=missing-function-docstring
+        if hasattr(self, "prefetch_data") and self.prefetch_data is not None:
+            self.prefetch_data.run_prefetcher = False
+            self.prefetch_data.stop_iteration = True
+            self.prefetch_data.paused = False
+            self.prefetch_data = None
+        if hasattr(self, "thread") and self.thread is not None:
+            self.thread.join()
+            self.thread = None
+
+    def __del__(self):
+        self.shutdown()
+
+    def __len__(self) -> int:
+        if isinstance(self.source_datapipe, Sized):
+            return len(self.source_datapipe)
+        raise TypeError(
+            f"{type(self).__name__} instance doesn't have valid length"
+        )

python/dgl/graphbolt/datapipes/visualization.py

@@ -0,0 +1,206 @@
+# pylint: disable=W,C,R
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# Original source:
+# https://github.com/pytorch/data/blob/v0.7.1/torchdata/datapipes/utils/_visualization.py
+
+import itertools
+from collections import defaultdict
+
+from typing import Optional, Set, TYPE_CHECKING
+
+from torch.utils.data.datapipes.iter.combining import _ChildDataPipe
+
+from .utils import IterDataPipe, traverse_dps
+
+if TYPE_CHECKING:
+    import graphviz
+
+
+__all__ = [
+    "to_graph",
+]
+
+
+class Node:
+    def __init__(self, dp, *, name=None):
+        self.dp = dp
+        self.name = name or type(dp).__name__.replace("IterDataPipe", "")
+        self.childs = set()
+        self.parents = set()
+
+    def add_child(self, child):
+        self.childs.add(child)
+        child.parents.add(self)
+
+    def remove_child(self, child):
+        self.childs.remove(child)
+        child.parents.remove(self)
+
+    def add_parent(self, parent):
+        self.parents.add(parent)
+        parent.childs.add(self)
+
+    def remove_parent(self, parent):
+        self.parents.remove(parent)
+        parent.childs.remove(self)
+
+    def __eq__(self, other):
+        if not isinstance(other, Node):
+            return NotImplemented
+
+        return hash(self) == hash(other)
+
+    def __hash__(self):
+        return hash(self.dp)
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return f"{self}-{hash(self)}"
+
+
+def to_nodes(dp, *, debug: bool) -> Set[Node]:
+    def recurse(dp_graph, child=None):
+        for _dp_id, (dp_node, dp_parents) in dp_graph.items():
+            node = Node(dp_node)
+            if child is not None:
+                node.add_child(child)
+            yield node
+            yield from recurse(dp_parents, child=node)
+
+    def aggregate(nodes):
+        groups = defaultdict(list)
+        for node in nodes:
+            groups[node].append(node)
+
+        nodes = set()
+        for node, group in groups.items():
+            if len(group) == 1:
+                nodes.add(node)
+                continue
+
+            aggregated_node = Node(node.dp)
+
+            for duplicate_node in group:
+                for child in duplicate_node.childs.copy():
+                    duplicate_node.remove_child(child)
+                    aggregated_node.add_child(child)
+
+                for parent in duplicate_node.parents.copy():
+                    duplicate_node.remove_parent(parent)
+                    aggregated_node.add_parent(parent)
+
+            nodes.add(aggregated_node)
+
+        if debug:
+            return nodes
+
+        child_dp_nodes = set(
+            itertools.chain.from_iterable(
+                node.parents
+                for node in nodes
+                if isinstance(node.dp, _ChildDataPipe)
+            )
+        )
+
+        if not child_dp_nodes:
+            return nodes
+
+        for node in child_dp_nodes:
+            fixed_parent_node = Node(
+                type(
+                    str(node).lstrip("_"),
+                    (IterDataPipe,),
+                    dict(dp=node.dp, childs=node.childs),
+                )()
+            )
+            nodes.remove(node)
+            nodes.add(fixed_parent_node)
+
+            for parent in node.parents.copy():
+                node.remove_parent(parent)
+                fixed_parent_node.add_parent(parent)
+
+            for child in node.childs:
+                nodes.remove(child)
+                for actual_child in child.childs.copy():
+                    actual_child.remove_parent(child)
+                    actual_child.add_parent(fixed_parent_node)
+
+        return nodes
+
+    return aggregate(recurse(traverse_dps(dp)))
+
+
+def to_graph(dp, *, debug: bool = False) -> "graphviz.Digraph":
+    """Visualizes a DataPipe by returning a :class:`graphviz.Digraph`, which is a graph of the data pipeline.
+    This allows you to visually inspect all the transformation that takes place in your DataPipes.
+
+    .. note::
+
+        The package :mod:`graphviz` is required to use this function.
+
+    .. note::
+
+        The most common interfaces for the returned graph object are:
+
+        - :meth:`~graphviz.Digraph.render`: Save the graph to a file.
+        - :meth:`~graphviz.Digraph.view`: Open the graph in a viewer.
+
+    Args:
+        dp: DataPipe that you would like to visualize (generally the last one in a chain of DataPipes).
+        debug (bool): If ``True``, renders internal datapipes that are usually hidden from the user
+            (such as ``ChildDataPipe`` of `demux` and `fork`). Defaults to ``False``.
+
+    Example:
+        >>> from torchdata.datapipes.iter import IterableWrapper
+        >>> from torchdata.datapipes.utils import to_graph
+        >>> dp = IterableWrapper(range(10))
+        >>> dp1, dp2 = dp.demux(num_instances=2, classifier_fn=lambda x: x % 2)
+        >>> dp1 = dp1.map(lambda x: x + 1)
+        >>> dp2 = dp2.filter(lambda _: True)
+        >>> dp3 = dp1.zip(dp2).map(lambda t: t[0] + t[1])
+        >>> g = to_graph(dp3)
+        >>> g.view()  # This will open the graph in a viewer
+    """
+    try:
+        import graphviz
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "The package `graphviz` is required to be installed to use this function. "
+            "Please `pip install graphviz` or `conda install -c conda-forge graphviz`."
+        ) from None
+
+    # The graph style as well as the color scheme below was copied from https://github.com/szagoruyko/pytorchviz/
+    # https://github.com/szagoruyko/pytorchviz/blob/0adcd83af8aa7ab36d6afd139cabbd9df598edb7/torchviz/dot.py#L78-L85
+    node_attr = dict(
+        style="filled",
+        shape="box",
+        align="left",
+        fontsize="10",
+        ranksep="0.1",
+        height="0.2",
+        fontname="monospace",
+    )
+    graph = graphviz.Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"))
+
+    for node in to_nodes(dp, debug=debug):
+        fillcolor: Optional[str]
+        if not node.parents:
+            fillcolor = "lightblue"
+        elif not node.childs:
+            fillcolor = "darkolivegreen1"
+        else:
+            fillcolor = None
+
+        graph.node(name=repr(node), label=str(node), fillcolor=fillcolor)
+
+        for child in node.childs:
+            graph.edge(repr(node), repr(child))
+
+    return graph

python/dgl/graphbolt/dataset.py

@@ -0,0 +1,95 @@
+"""GraphBolt Dataset."""
+
+from typing import Dict, List, Union
+
+from .feature_store import FeatureStore
+from .itemset import HeteroItemSet, ItemSet
+from .sampling_graph import SamplingGraph
+
+__all__ = [
+    "Task",
+    "Dataset",
+]
+
+
+class Task:
+    """An abstract task which consists of meta information and
+    Train/Validation/Test Set.
+
+    * meta information
+        The meta information of a task includes any kinds of data that are
+        defined by the user in YAML when instantiating the task.
+
+    * Train/Validation/Test Set
+        The train/validation/test (TVT) set which is used to train the neural
+        networks. We calculate the embeddings based on their respective features
+        and the graph structure, and then utilize the embeddings to optimize the
+        neural network parameters.
+    """
+
+    @property
+    def metadata(self) -> Dict:
+        """Return the task metadata."""
+        raise NotImplementedError
+
+    @property
+    def train_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the training set."""
+        raise NotImplementedError
+
+    @property
+    def validation_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the validation set."""
+        raise NotImplementedError
+
+    @property
+    def test_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the test set."""
+        raise NotImplementedError
+
+
+class Dataset:
+    """An abstract dataset which provides abstraction for accessing the data
+    required for training.
+
+    The data abstraction could be a native CPU memory block, a shared memory
+    block, a file handle of an opened file on disk, a service that provides
+    the API to access the data e.t.c. There are 3 primary components in the
+    dataset:
+
+    * Task
+        A task consists of several meta information and the
+        Train/Validation/Test Set. A dataset could have multiple tasks.
+
+    * Feature Storage
+        A key-value store which stores node/edge/graph features.
+
+    * Graph Topology
+        Graph topology is used by the subgraph sampling algorithm to generate
+        a subgraph.
+    """
+
+    @property
+    def tasks(self) -> List[Task]:
+        """Return the tasks."""
+        raise NotImplementedError
+
+    @property
+    def graph(self) -> SamplingGraph:
+        """Return the graph."""
+        raise NotImplementedError
+
+    @property
+    def feature(self) -> FeatureStore:
+        """Return the feature."""
+        raise NotImplementedError
+
+    @property
+    def dataset_name(self) -> str:
+        """Return the dataset name."""
+        raise NotImplementedError
+
+    @property
+    def all_nodes_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the itemset containing all nodes."""
+        raise NotImplementedError

python/dgl/graphbolt/external_utils.py

@@ -0,0 +1,159 @@
+"""Utility functions for external use."""
+from functools import partial
+from typing import Dict, Union
+
+import torch
+
+from torch.utils.data import functional_datapipe
+
+from .minibatch import MiniBatch
+from .minibatch_transformer import MiniBatchTransformer
+
+
+@functional_datapipe("exclude_seed_edges")
+class SeedEdgesExcluder(MiniBatchTransformer):
+    """A mini-batch transformer used to manipulate mini-batch.
+
+    Functional name: :obj:`transform`.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    include_reverse_edges : bool
+        Whether reverse edges should be excluded as well. Default is False.
+    reverse_etypes_mapping : Dict[str, str] = None
+        The mapping from the original edge types to their reverse edge types.
+    asynchronous: bool
+        Boolean indicating whether edge exclusion stages should run on
+        background threads to hide the latency of CPU GPU synchronization.
+        Should be enabled only when sampling on the GPU.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        include_reverse_edges: bool = False,
+        reverse_etypes_mapping: Dict[str, str] = None,
+        asynchronous=False,
+    ):
+        exclude_seed_edges_fn = partial(
+            exclude_seed_edges,
+            include_reverse_edges=include_reverse_edges,
+            reverse_etypes_mapping=reverse_etypes_mapping,
+            async_op=asynchronous,
+        )
+        datapipe = datapipe.transform(exclude_seed_edges_fn)
+        if asynchronous:
+            datapipe = datapipe.buffer()
+            datapipe = datapipe.transform(self._wait_for_sampled_subgraphs)
+        super().__init__(datapipe)
+
+    @staticmethod
+    def _wait_for_sampled_subgraphs(minibatch):
+        minibatch.sampled_subgraphs = [
+            subgraph.wait() for subgraph in minibatch.sampled_subgraphs
+        ]
+        return minibatch
+
+
+def add_reverse_edges(
+    edges: Union[Dict[str, torch.Tensor], torch.Tensor],
+    reverse_etypes_mapping: Dict[str, str] = None,
+):
+    r"""
+    This function finds the reverse edges of the given `edges` and returns the
+    composition of them. In a homogeneous graph, reverse edges have inverted
+    source and destination node IDs. While in a heterogeneous graph, reversing
+    also involves swapping node IDs and their types. This function could be
+    used before `exclude_edges` function to help find targeting edges.
+    Note: The found reverse edges may not really exists in the original graph.
+    And repeat edges could be added becasue reverse edges may already exists in
+    the `edges`.
+
+    Parameters
+    ----------
+    edges : Union[Dict[str, torch.Tensor], torch.Tensor]
+      - If sampled subgraph is homogeneous, then `edges` should be a N*2
+        tensors.
+      - If sampled subgraph is heterogeneous, then `edges` should be a
+        dictionary of edge types and the corresponding edges to exclude.
+    reverse_etypes_mapping : Dict[str, str], optional
+        The mapping from the original edge types to their reverse edge types.
+
+    Returns
+    -------
+    Union[Dict[str, torch.Tensor], torch.Tensor]
+        The node pairs contain both the original edges and their reverse
+        counterparts.
+
+    Examples
+    --------
+    >>> edges = {"A:r:B": torch.tensor([[0, 1],[1, 2]]))}
+    >>> print(gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A"}))
+    {'A:r:B': torch.tensor([[0, 1],[1, 2]]),
+    'B:rr:A': torch.tensor([[1, 0],[2, 1]])}
+
+    >>> edges = torch.tensor([[0, 1],[1, 2]])
+    >>> print(gb.add_reverse_edges(edges))
+    torch.tensor([[1, 0],[2, 1]])
+    """
+    if isinstance(edges, torch.Tensor):
+        assert edges.ndim == 2 and edges.shape[1] == 2, (
+            "Only tensor with shape N*2 is supported now, but got "
+            + f"{edges.shape}."
+        )
+        reverse_edges = edges.flip(dims=(1,))
+        return torch.cat((edges, reverse_edges))
+    else:
+        combined_edges = edges.copy()
+        for etype, reverse_etype in reverse_etypes_mapping.items():
+            if etype in edges:
+                assert edges[etype].ndim == 2 and edges[etype].shape[1] == 2, (
+                    "Only tensor with shape N*2 is supported now, but got "
+                    + f"{edges[etype].shape}."
+                )
+                if reverse_etype in combined_edges:
+                    combined_edges[reverse_etype] = torch.cat(
+                        (
+                            combined_edges[reverse_etype],
+                            edges[etype].flip(dims=(1,)),
+                        )
+                    )
+                else:
+                    combined_edges[reverse_etype] = edges[etype].flip(dims=(1,))
+        return combined_edges
+
+
+def exclude_seed_edges(
+    minibatch: MiniBatch,
+    include_reverse_edges: bool = False,
+    reverse_etypes_mapping: Dict[str, str] = None,
+    async_op: bool = False,
+):
+    """
+    Exclude seed edges with or without their reverse edges from the sampled
+    subgraphs in the minibatch.
+
+    Parameters
+    ----------
+    minibatch : MiniBatch
+        The minibatch.
+    include_reverse_edges : bool
+        Whether reverse edges should be excluded as well. Default is False.
+    reverse_etypes_mapping : Dict[str, str] = None
+        The mapping from the original edge types to their reverse edge types.
+    async_op: bool
+        Boolean indicating whether the call is asynchronous. If so, the result
+        can be obtained by calling wait on the modified sampled_subgraphs.
+    """
+    edges_to_exclude = minibatch.seeds
+    if include_reverse_edges:
+        edges_to_exclude = add_reverse_edges(
+            edges_to_exclude, reverse_etypes_mapping
+        )
+    minibatch.sampled_subgraphs = [
+        subgraph.exclude_edges(edges_to_exclude, async_op=async_op)
+        for subgraph in minibatch.sampled_subgraphs
+    ]
+    return minibatch

python/dgl/graphbolt/feature_fetcher.py

@@ -0,0 +1,282 @@
+"""Feature fetchers"""
+
+from functools import partial
+from typing import Dict
+
+import torch
+
+from torch.utils.data import functional_datapipe
+
+from .base import etype_tuple_to_str
+from .impl.cooperative_conv import CooperativeConvFunction
+
+from .minibatch_transformer import MiniBatchTransformer
+
+
+__all__ = [
+    "FeatureFetcher",
+    "FeatureFetcherStartMarker",
+]
+
+
+def get_feature_key_list(feature_keys, domain):
+    """Processes node_feature_keys and extracts their feature keys to a list."""
+    if isinstance(feature_keys, Dict):
+        return [
+            (domain, type_name, feature_name)
+            for type_name, feature_names in feature_keys.items()
+            for feature_name in feature_names
+        ]
+    elif feature_keys is not None:
+        return [(domain, None, feature_name) for feature_name in feature_keys]
+    else:
+        return []
+
+
+@functional_datapipe("mark_feature_fetcher_start")
+class FeatureFetcherStartMarker(MiniBatchTransformer):
+    """Used to mark the start of a FeatureFetcher and is a no-op. All the
+    datapipes created during a FeatureFetcher instantiation are guarenteed to be
+    contained between FeatureFetcherStartMarker and FeatureFetcher instances in
+    the datapipe graph.
+    """
+
+    def __init__(self, datapipe):
+        super().__init__(datapipe, self._identity)
+
+
+@functional_datapipe("fetch_feature")
+class FeatureFetcher(MiniBatchTransformer):
+    """A feature fetcher used to fetch features for node/edge in graphbolt.
+
+    Functional name: :obj:`fetch_feature`.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    feature_store : FeatureStore
+        A storage for features, support read and update.
+    node_feature_keys : List[str] or Dict[str, List[str]]
+        Node features keys indicates the node features need to be read.
+        - If `node_features` is a list: It means the graph is homogeneous
+        graph, and the 'str' inside are feature names.
+        - If `node_features` is a dictionary: The keys should be node type
+        and the values are lists of feature names.
+    edge_feature_keys : List[str] or Dict[str, List[str]]
+        Edge features name indicates the edge features need to be read.
+        - If `edge_features` is a list: It means the graph is homogeneous
+        graph, and the 'str' inside are feature names.
+        - If `edge_features` is a dictionary: The keys are edge types,
+        following the format 'str:str:str', and the values are lists of
+        feature names.
+    overlap_fetch : bool, optional
+        If True, the feature fetcher will overlap the UVA feature fetcher
+        operations with the rest of operations by using an alternative CUDA
+        stream or utilizing asynchronous operations. Default is True.
+    cooperative: bool, optional
+        Boolean indicating whether Cooperative Minibatching, which was initially
+        proposed in
+        `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__
+        and was later first fully described in
+        `Cooperative Minibatching in Graph Neural Networks
+        <https://arxiv.org/abs/2310.12403>`__. Cooperation between the GPUs
+        eliminates duplicate work performed across the GPUs due to the
+        overlapping sampled k-hop neighborhoods of seed nodes when performing
+        GNN minibatching.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        feature_store,
+        node_feature_keys=None,
+        edge_feature_keys=None,
+        overlap_fetch=True,
+        cooperative=False,
+    ):
+        datapipe = datapipe.mark_feature_fetcher_start()
+        self.feature_store = feature_store
+        self.node_feature_keys = node_feature_keys
+        self.edge_feature_keys = edge_feature_keys
+        max_val = 0
+        if overlap_fetch:
+            for feature_key_list in [
+                get_feature_key_list(node_feature_keys, "node"),
+                get_feature_key_list(edge_feature_keys, "edge"),
+            ]:
+                for feature_key in feature_key_list:
+                    if feature_key not in feature_store:
+                        continue
+                    for device_str in ["cpu", "cuda"]:
+                        try:
+                            max_val = max(
+                                feature_store[
+                                    feature_key
+                                ].read_async_num_stages(
+                                    torch.device(device_str)
+                                ),
+                                max_val,
+                            )
+                        except AssertionError:
+                            pass
+        datapipe = datapipe.transform(self._read)
+        for i in range(max_val, 0, -1):
+            datapipe = datapipe.transform(
+                partial(self._execute_stage, i)
+            ).buffer(1)
+        if max_val > 0:
+            datapipe = datapipe.transform(self._final_stage)
+        if cooperative:
+            datapipe = datapipe.transform(self._cooperative_exchange)
+            datapipe = datapipe.buffer()
+        super().__init__(datapipe)
+        # A positive value indicates that the overlap optimization is enabled.
+        self.max_num_stages = max_val
+
+    @staticmethod
+    def _execute_stage(current_stage, data):
+        all_features = [data.node_features] + [
+            data.edge_features[i] for i in range(data.num_layers())
+        ]
+        for features in all_features:
+            for key in features:
+                handle, stage = features[key]
+                assert current_stage >= stage
+                if current_stage == stage:
+                    value = next(handle)
+                    features[key] = (handle if stage > 1 else value, stage - 1)
+        return data
+
+    @staticmethod
+    def _final_stage(data):
+        all_features = [data.node_features] + [
+            data.edge_features[i] for i in range(data.num_layers())
+        ]
+        for features in all_features:
+            for key in features:
+                value, stage = features[key]
+                assert stage == 0
+                features[key] = value.wait()
+        return data
+
+    def _cooperative_exchange(self, data):
+        subgraph = data.sampled_subgraphs[0]
+        is_heterogeneous = isinstance(
+            self.node_feature_keys, Dict
+        ) or isinstance(self.edge_feature_keys, Dict)
+        if is_heterogeneous:
+            node_features = {key: {} for key, _ in data.node_features.keys()}
+            for (key, ntype), feature in data.node_features.items():
+                node_features[key][ntype] = feature
+            for key, feature in node_features.items():
+                new_feature = CooperativeConvFunction.apply(subgraph, feature)
+                for ntype, tensor in new_feature.items():
+                    data.node_features[(key, ntype)] = tensor
+        else:
+            for key in data.node_features:
+                feature = data.node_features[key]
+                new_feature = CooperativeConvFunction.apply(subgraph, feature)
+                data.node_features[key] = new_feature
+        return data
+
+    def _read(self, data):
+        """
+        Fill in the node/edge features field in data.
+
+        Parameters
+        ----------
+        data : MiniBatch
+            An instance of :class:`MiniBatch`. Even if 'node_feature' or
+            'edge_feature' is already filled, it will be overwritten for
+            overlapping features.
+
+        Returns
+        -------
+        MiniBatch
+            An instance of :class:`MiniBatch` filled with required features.
+        """
+        node_features = {}
+        num_layers = data.num_layers()
+        edge_features = [{} for _ in range(num_layers)]
+        is_heterogeneous = isinstance(
+            self.node_feature_keys, Dict
+        ) or isinstance(self.edge_feature_keys, Dict)
+        # Read Node features.
+        input_nodes = data.node_ids()
+
+        def read_helper(feature_key, index):
+            if self.max_num_stages > 0:
+                feature = self.feature_store[feature_key]
+                num_stages = feature.read_async_num_stages(index.device)
+                if num_stages > 0:
+                    return (feature.read_async(index), num_stages)
+                else:  # Asynchronicity is not needed, compute in _final_stage.
+
+                    class _Waiter:
+                        def __init__(self, feature, index):
+                            self.feature = feature
+                            self.index = index
+
+                        def wait(self):
+                            """Returns the stored value when invoked."""
+                            result = self.feature.read(self.index)
+                            # Ensure there is no memory leak.
+                            self.feature = self.index = None
+                            return result
+
+                    return (_Waiter(feature, index), 0)
+            else:
+                domain, type_name, feature_name = feature_key
+                return self.feature_store.read(
+                    domain, type_name, feature_name, index
+                )
+
+        if self.node_feature_keys and input_nodes is not None:
+            if is_heterogeneous:
+                for type_name, nodes in input_nodes.items():
+                    if type_name not in self.node_feature_keys or nodes is None:
+                        continue
+                    for feature_name in self.node_feature_keys[type_name]:
+                        node_features[(type_name, feature_name)] = read_helper(
+                            ("node", type_name, feature_name), nodes
+                        )
+            else:
+                for feature_name in self.node_feature_keys:
+                    node_features[feature_name] = read_helper(
+                        ("node", None, feature_name), input_nodes
+                    )
+        # Read Edge features.
+        if self.edge_feature_keys and num_layers > 0:
+            for i in range(num_layers):
+                original_edge_ids = data.edge_ids(i)
+                if is_heterogeneous:
+                    # Convert edge type to string.
+                    original_edge_ids = {
+                        (
+                            etype_tuple_to_str(key)
+                            if isinstance(key, tuple)
+                            else key
+                        ): value
+                        for key, value in original_edge_ids.items()
+                    }
+                    for type_name, edges in original_edge_ids.items():
+                        if (
+                            type_name not in self.edge_feature_keys
+                            or edges is None
+                        ):
+                            continue
+                        for feature_name in self.edge_feature_keys[type_name]:
+                            edge_features[i][
+                                (type_name, feature_name)
+                            ] = read_helper(
+                                ("edge", type_name, feature_name), edges
+                            )
+                else:
+                    for feature_name in self.edge_feature_keys:
+                        edge_features[i][feature_name] = read_helper(
+                            ("edge", None, feature_name), original_edge_ids
+                        )
+        data.set_node_features(node_features)
+        data.set_edge_features(edge_features)
+        return data

python/dgl/graphbolt/feature_store.py

@@ -0,0 +1,351 @@
+"""Feature store for GraphBolt."""
+
+from typing import Dict, NamedTuple, Union
+
+import torch
+
+__all__ = [
+    "bytes_to_number_of_items",
+    "Feature",
+    "FeatureStore",
+    "FeatureKey",
+    "wrap_with_cached_feature",
+]
+
+
+class FeatureKey(NamedTuple):
+    """A named tuple class to represent feature keys in FeatureStore classes.
+    The fields are domain, type and name all of which take string values.
+    """
+
+    domain: str
+    type: str
+    name: int
+
+
+class Feature:
+    r"""A wrapper of feature data for access."""
+
+    def __init__(self):
+        pass
+
+    def read(self, ids: torch.Tensor = None):
+        """Read from the feature.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        raise NotImplementedError
+
+    def read_async(self, ids: torch.Tensor):
+        """Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            `read_async_num_stages(ids.device)`th invocation. The return result
+            can be accessed by calling `.wait()`. on the returned future object.
+            It is undefined behavior to call `.wait()` more than once.
+
+        Example Usage
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        raise NotImplementedError
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        raise NotImplementedError
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        raise NotImplementedError
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        raise NotImplementedError
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        raise NotImplementedError
+
+    def metadata(self):
+        """Get the metadata of the feature.
+
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return {}
+
+
+class FeatureStore:
+    r"""A store to manage multiple features for access."""
+
+    def __init__(self):
+        pass
+
+    def __getitem__(self, feature_key: FeatureKey) -> Feature:
+        """Access the underlying `Feature` with its (domain, type, name) as
+        the feature_key.
+        """
+        raise NotImplementedError
+
+    def __setitem__(self, feature_key: FeatureKey, feature: Feature):
+        """Set the underlying `Feature` with its (domain, type, name) as
+        the feature_key and feature as the value.
+        """
+        raise NotImplementedError
+
+    def __contains__(self, feature_key: FeatureKey) -> bool:
+        """Checks whether the provided (domain, type, name) as the feature_key
+        is container in the FeatureStore."""
+        raise NotImplementedError
+
+    def read(
+        self,
+        domain: str,
+        type_name: str,
+        feature_name: str,
+        ids: torch.Tensor = None,
+    ):
+        """Read from the feature store.
+
+        Parameters
+        ----------
+        domain : str
+            The domain of the feature such as "node", "edge" or "graph".
+        type_name : str
+            The node or edge type name.
+        feature_name : str
+            The feature name.
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        return self.__getitem__((domain, type_name, feature_name)).read(ids)
+
+    def size(
+        self,
+        domain: str,
+        type_name: str,
+        feature_name: str,
+    ):
+        """Get the size of the specified feature in the feature store.
+
+        Parameters
+        ----------
+        domain : str
+            The domain of the feature such as "node", "edge" or "graph".
+        type_name : str
+            The node or edge type name.
+        feature_name : str
+            The feature name.
+        Returns
+        -------
+        torch.Size
+            The size of the specified feature in the feature store.
+        """
+        return self.__getitem__((domain, type_name, feature_name)).size()
+
+    def count(
+        self,
+        domain: str,
+        type_name: str,
+        feature_name: str,
+    ):
+        """Get the count the specified feature in the feature store.
+
+        Parameters
+        ----------
+        domain : str
+            The domain of the feature such as "node", "edge" or "graph".
+        type_name : str
+            The node or edge type name.
+        feature_name : str
+            The feature name.
+        Returns
+        -------
+        int
+            The count of the specified feature in the feature store.
+        """
+        return self.__getitem__((domain, type_name, feature_name)).count()
+
+    def metadata(
+        self,
+        domain: str,
+        type_name: str,
+        feature_name: str,
+    ):
+        """Get the metadata of the specified feature in the feature store.
+
+        Parameters
+        ----------
+        domain : str
+            The domain of the feature such as "node", "edge" or "graph".
+        type_name : str
+            The node or edge type name.
+        feature_name : str
+            The feature name.
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return self.__getitem__((domain, type_name, feature_name)).metadata()
+
+    def update(
+        self,
+        domain: str,
+        type_name: str,
+        feature_name: str,
+        value: torch.Tensor,
+        ids: torch.Tensor = None,
+    ):
+        """Update the feature store.
+
+        Parameters
+        ----------
+        domain : str
+            The domain of the feature such as "node", "edge" or "graph".
+        type_name : str
+            The node or edge type name.
+        feature_name : str
+            The feature name.
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        self.__getitem__((domain, type_name, feature_name)).update(value, ids)
+
+    def keys(self):
+        """Get the keys of the features.
+
+        Returns
+        -------
+        List[tuple]
+            The keys of the features. The tuples are in `(domain, type_name,
+            feat_name)` format.
+        """
+        raise NotImplementedError
+
+
+def bytes_to_number_of_items(cache_capacity_in_bytes, single_item):
+    """Returns the number of rows to be cached."""
+    item_bytes = single_item.nbytes
+    # Round up so that we never get a size of 0, unless bytes is 0.
+    return (cache_capacity_in_bytes + item_bytes - 1) // item_bytes
+
+
+def wrap_with_cached_feature(
+    cached_feature_type,
+    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
+    max_cache_size_in_bytes: int,
+    *args,
+    **kwargs,
+) -> Union[Feature, Dict[FeatureKey, Feature]]:
+    """Wraps the given features with the given cached feature type using
+    a single cache instance."""
+    if not isinstance(fallback_features, dict):
+        assert isinstance(fallback_features, Feature)
+        return wrap_with_cached_feature(
+            cached_feature_type,
+            {"a": fallback_features},
+            max_cache_size_in_bytes,
+            *args,
+            **kwargs,
+        )["a"]
+    row_bytes = None
+    cache = None
+    wrapped_features = {}
+    offset = 0
+    for feature_key, fallback_feature in fallback_features.items():
+        # Fetching the feature dimension from the underlying feature.
+        feat0 = fallback_feature.read(torch.tensor([0]))
+        if row_bytes is None:
+            row_bytes = feat0.nbytes
+        else:
+            assert (
+                row_bytes == feat0.nbytes
+            ), "The # bytes of a single row of the features should match."
+        cache_size = bytes_to_number_of_items(max_cache_size_in_bytes, feat0)
+        if cache is None:
+            cache = cached_feature_type._cache_type(
+                cache_shape=(cache_size,) + feat0.shape[1:],
+                dtype=feat0.dtype,
+                *args,
+                **kwargs,
+            )
+        wrapped_features[feature_key] = cached_feature_type(
+            fallback_feature, cache=cache, offset=offset
+        )
+        offset += fallback_feature.count()
+
+    return wrapped_features

python/dgl/graphbolt/impl/__init__.py

@@ -0,0 +1,18 @@
+"""Implementation of GraphBolt."""
+from .basic_feature_store import *
+from .fused_csc_sampling_graph import *
+from .gpu_feature_cache import *
+from .gpu_cached_feature import *
+from .in_subgraph_sampler import *
+from .legacy_dataset import *
+from .neighbor_sampler import *
+from .temporal_neighbor_sampler import *
+from .ondisk_dataset import *
+from .ondisk_metadata import *
+from .sampled_subgraph_impl import *
+from .torch_based_feature_store import *
+from .uniform_negative_sampler import *
+from .gpu_graph_cache import *
+from .cpu_feature_cache import *
+from .cpu_cached_feature import *
+from .cooperative_conv import *

python/dgl/graphbolt/impl/basic_feature_store.py

@@ -0,0 +1,60 @@
+"""Basic feature store for GraphBolt."""
+
+from typing import Dict, Tuple
+
+from ..feature_store import Feature, FeatureKey, FeatureStore
+
+__all__ = ["BasicFeatureStore"]
+
+
+class BasicFeatureStore(FeatureStore):
+    r"""A basic feature store to manage multiple features for access."""
+
+    def __init__(self, features: Dict[Tuple[str, str, str], Feature]):
+        r"""Initiate a basic feature store.
+
+
+        Parameters
+        ----------
+        features : Dict[Tuple[str, str, str], Feature]
+            The dict of features served by the feature store, in which the key
+            is tuple of (domain, type_name, feature_name).
+
+        Returns
+        -------
+        The feature stores.
+        """
+        super().__init__()
+        self._features = features
+
+    def __getitem__(self, feature_key: FeatureKey) -> Feature:
+        """Access the underlying `Feature` with its (domain, type, name) as
+        the feature_key.
+        """
+        return self._features[feature_key]
+
+    def __setitem__(self, feature_key: FeatureKey, feature: Feature):
+        """Set the underlying `Feature` with its (domain, type, name) as
+        the feature_key and feature as the value.
+        """
+        self._features[feature_key] = feature
+
+    def __contains__(self, feature_key: FeatureKey) -> bool:
+        """Checks whether the provided (domain, type, name) as the feature_key
+        is container in the BasicFeatureStore."""
+        return feature_key in self._features
+
+    def __len__(self):
+        """Return the number of features."""
+        return len(self._features)
+
+    def keys(self):
+        """Get the keys of the features.
+
+        Returns
+        -------
+        List[tuple]
+            The keys of the features. The tuples are in `(domain, type_name,
+            feat_name)` format.
+        """
+        return list(self._features.keys())

python/dgl/graphbolt/impl/cooperative_conv.py

@@ -0,0 +1,117 @@
+"""Graphbolt cooperative convolution."""
+from typing import Dict, Union
+
+import torch
+
+from ..sampled_subgraph import SampledSubgraph
+from ..subgraph_sampler import all_to_all, convert_to_hetero, revert_to_homo
+
+__all__ = ["CooperativeConvFunction", "CooperativeConv"]
+
+
+class CooperativeConvFunction(torch.autograd.Function):
+    """Cooperative convolution operation from Cooperative Minibatching.
+
+    Implements the `all-to-all` message passing algorithm
+    in Cooperative Minibatching, which was initially proposed in
+    `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
+    was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching. This reduces the redundant computations across
+    GPUs at the expense of communication.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        subgraph: SampledSubgraph,
+        tensor: Union[torch.Tensor, Dict[str, torch.Tensor]],
+    ):
+        """Implements the forward pass."""
+        counts_sent = convert_to_hetero(subgraph._counts_sent)
+        counts_received = convert_to_hetero(subgraph._counts_received)
+        seed_inverse_ids = convert_to_hetero(subgraph._seed_inverse_ids)
+        seed_sizes = convert_to_hetero(subgraph._seed_sizes)
+        ctx.communication_variables = (
+            counts_sent,
+            counts_received,
+            seed_inverse_ids,
+            seed_sizes,
+        )
+        outs = {}
+        for ntype, typed_tensor in convert_to_hetero(tensor).items():
+            out = typed_tensor.new_empty(
+                (sum(counts_sent[ntype]),) + typed_tensor.shape[1:]
+            )
+            all_to_all(
+                torch.split(out, counts_sent[ntype]),
+                torch.split(
+                    typed_tensor[seed_inverse_ids[ntype]],
+                    counts_received[ntype],
+                ),
+            )
+            outs[ntype] = out
+        return revert_to_homo(out)
+
+    @staticmethod
+    def backward(
+        ctx, grad_output: Union[torch.Tensor, Dict[str, torch.Tensor]]
+    ):
+        """Implements the backward pass."""
+        (
+            counts_sent,
+            counts_received,
+            seed_inverse_ids,
+            seed_sizes,
+        ) = ctx.communication_variables
+        delattr(ctx, "communication_variables")
+        outs = {}
+        for ntype, typed_grad_output in convert_to_hetero(grad_output).items():
+            out = typed_grad_output.new_empty(
+                (sum(counts_received[ntype]),) + typed_grad_output.shape[1:]
+            )
+            all_to_all(
+                torch.split(out, counts_received[ntype]),
+                torch.split(typed_grad_output, counts_sent[ntype]),
+            )
+            i = out.new_empty(2, out.shape[0], dtype=torch.int64)
+            i[0] = seed_inverse_ids[ntype]  # src
+            i[1] = torch.arange(
+                out.shape[0], device=typed_grad_output.device
+            )  # dst
+            coo = torch.sparse_coo_tensor(
+                i,
+                torch.ones(
+                    i.shape[1], dtype=grad_output.dtype, device=i.device
+                ),
+                size=(seed_sizes[ntype], i.shape[1]),
+            )
+            outs[ntype] = torch.sparse.mm(coo, out)
+        return None, revert_to_homo(outs)
+
+
+class CooperativeConv(torch.nn.Module):
+    """Cooperative convolution operation from Cooperative Minibatching.
+
+    Implements the `all-to-all` message passing algorithm
+    in Cooperative Minibatching, which was initially proposed in
+    `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
+    was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching. This reduces the redundant computations across
+    GPUs at the expense of communication.
+    """
+
+    def forward(
+        self,
+        subgraph: SampledSubgraph,
+        x: Union[torch.Tensor, Dict[str, torch.Tensor]],
+    ):
+        """Implements the forward pass."""
+        return CooperativeConvFunction.apply(subgraph, x)

python/dgl/graphbolt/impl/cpu_cached_feature.py

@@ -0,0 +1,499 @@
+"""CPU cached feature for GraphBolt."""
+from typing import Dict, Optional, Union
+
+import torch
+
+from ..base import get_device_to_host_uva_stream, get_host_to_device_uva_stream
+from ..feature_store import (
+    bytes_to_number_of_items,
+    Feature,
+    FeatureKey,
+    wrap_with_cached_feature,
+)
+
+from .cpu_feature_cache import CPUFeatureCache
+
+__all__ = ["CPUCachedFeature", "cpu_cached_feature"]
+
+
+class CPUCachedFeature(Feature):
+    r"""CPU cached feature wrapping a fallback feature. Use `cpu_cached_feature`
+    to construct an instance of this class.
+
+    Parameters
+    ----------
+    fallback_feature : Feature
+        The fallback feature.
+    cache : CPUFeatureCache
+        A CPUFeatureCache instance to serve as the cache backend.
+    offset : int, optional
+        The offset value to add to the given ids before using the cache. This
+        parameter is useful if multiple `CPUCachedFeature`s are sharing a single
+        CPUFeatureCache object.
+    """
+
+    _cache_type = CPUFeatureCache
+
+    def __init__(
+        self,
+        fallback_feature: Feature,
+        cache: CPUFeatureCache,
+        offset: int = 0,
+    ):
+        super(CPUCachedFeature, self).__init__()
+        assert isinstance(fallback_feature, Feature), (
+            f"The fallback_feature must be an instance of Feature, but got "
+            f"{type(fallback_feature)}."
+        )
+        self._fallback_feature = fallback_feature
+        self._feature = cache
+        self._offset = offset
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._fallback_feature.read()
+        return self._feature.query_and_replace(
+            ids.cpu(), self._fallback_feature.read, self._offset
+        ).to(ids.device)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        policy = self._feature._policy
+        cache = self._feature._cache
+        if ids.is_cuda and self.is_pinned():
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to("cpu", non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                found_positions = found_positions.to(
+                    ids_device, non_blocking=True
+                )
+                values_from_cpu = cache.index_select(found_positions)
+                values_from_cpu.record_stream(current_stream)
+                values_from_cpu_copy_event = torch.cuda.Event()
+                values_from_cpu_copy_event.record()
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values_from_cpu_copy_event.synchronize()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                index = index.to(ids_device, non_blocking=True)
+                missing_values = missing_values.to(
+                    ids_device, non_blocking=True
+                )
+                index.record_stream(current_stream)
+                missing_values.record_stream(current_stream)
+                missing_values_copy_event = torch.cuda.Event()
+                missing_values_copy_event.record()
+
+            yield
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, events, existing, missing, index):
+                    self.events = events
+                    self.existing = existing
+                    self.missing = missing
+                    self.index = index
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    for event in self.events:
+                        event.wait()
+                    values = torch.empty(
+                        (self.index.shape[0],) + self.missing.shape[1:],
+                        dtype=self.missing.dtype,
+                        device=ids_device,
+                    )
+                    num_found = self.existing.size(0)
+                    found_index = self.index[:num_found]
+                    missing_index = self.index[num_found:]
+                    values[found_index] = self.existing
+                    values[missing_index] = self.missing
+                    # Ensure there is no memory leak.
+                    self.events = self.existing = None
+                    self.missing = self.index = None
+                    return values
+
+            yield _Waiter(
+                [
+                    writing_completed,
+                    values_from_cpu_copy_event,
+                    missing_values_copy_event,
+                ],
+                values_from_cpu,
+                missing_values,
+                index,
+            )
+        elif ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to("cpu", non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            values_future = cache.query_async(
+                found_positions, index, ids.shape[0]
+            )
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values = values_future.wait()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_index = index[found_cnt:]
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+            values = torch.ops.graphbolt.scatter_async(
+                values, missing_index, missing_values
+            )
+
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values = values.wait().to(ids_device, non_blocking=True)
+                values.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, events, values):
+                    self.events = events
+                    self.values = values
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    for event in self.events:
+                        event.wait()
+                    values = self.values
+                    # Ensure there is no memory leak.
+                    self.events = self.values = None
+                    return values
+
+            yield _Waiter([values_copy_event, writing_completed], values)
+        else:
+            policy_future = policy.query_and_replace_async(ids, self._offset)
+
+            yield
+
+            (
+                positions,
+                index,
+                pointers,
+                missing_keys,
+                found_offsets,
+                missing_offsets,
+            ) = policy_future.wait()
+            self._feature.total_queries += ids.shape[0]
+            self._feature.total_miss += missing_keys.shape[0]
+            found_cnt = ids.size(0) - missing_keys.size(0)
+            found_positions = positions[:found_cnt]
+            missing_positions = positions[found_cnt:]
+            found_pointers = pointers[:found_cnt]
+            missing_pointers = pointers[found_cnt:]
+            values_future = cache.query_async(
+                found_positions, index, ids.shape[0]
+            )
+
+            fallback_reader = self._fallback_feature.read_async(missing_keys)
+            for _ in range(
+                self._fallback_feature.read_async_num_stages(
+                    missing_keys.device
+                )
+            ):
+                missing_values_future = next(fallback_reader, None)
+                yield  # fallback feature stages.
+
+            values = values_future.wait()
+            reading_completed = policy.reading_completed_async(
+                found_pointers, found_offsets
+            )
+
+            missing_index = index[found_cnt:]
+
+            missing_values = missing_values_future.wait()
+            replace_future = cache.replace_async(
+                missing_positions, missing_values
+            )
+            values = torch.ops.graphbolt.scatter_async(
+                values, missing_index, missing_values
+            )
+
+            yield
+
+            reading_completed.wait()
+            replace_future.wait()
+            writing_completed = policy.writing_completed_async(
+                missing_pointers, missing_offsets
+            )
+
+            class _Waiter:
+                def __init__(self, event, values):
+                    self.event = event
+                    self.values = values
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    self.event.wait()
+                    values = self.values.wait()
+                    # Ensure there is no memory leak.
+                    self.event = self.values = None
+                    return values
+
+            yield _Waiter(writing_completed, values)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        if ids_device.type == "cuda":
+            return 4 + self._fallback_feature.read_async_num_stages(
+                torch.device("cpu")
+            )
+        else:
+            return 3 + self._fallback_feature.read_async_num_stages(ids_device)
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._fallback_feature.size()
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._fallback_feature.count()
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            feat0 = value[:1]
+            self._fallback_feature.update(value)
+            cache_size = min(
+                bytes_to_number_of_items(self.cache_size_in_bytes, feat0),
+                value.shape[0],
+            )
+            self._feature = None  # Destroy the existing cache first.
+            self._feature = self._cache_type(
+                (cache_size,) + feat0.shape[1:], feat0.dtype
+            )
+        else:
+            self._fallback_feature.update(value, ids)
+            self._feature.replace(ids, value, None, self._offset)
+
+    def is_pinned(self):
+        """Returns True if the cache storage is pinned."""
+        return self._feature.is_pinned()
+
+    @property
+    def cache_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._feature.max_size_in_bytes
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self._feature.miss_rate
+
+
+def cpu_cached_feature(
+    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
+    max_cache_size_in_bytes: int,
+    policy: Optional[str] = None,
+    pin_memory: bool = False,
+) -> Union[CPUCachedFeature, Dict[FeatureKey, CPUCachedFeature]]:
+    r"""CPU cached feature wrapping a fallback feature.
+
+    Parameters
+    ----------
+    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
+        The fallback feature(s).
+    max_cache_size_in_bytes : int
+        The capacity of the cache in bytes. The size should be a few factors
+        larger than the size of each read request. Otherwise, the caching policy
+        will hang due to all cache entries being read and/or write locked,
+        resulting in a deadlock.
+    policy : str, optional
+        The cache eviction policy algorithm name. The available policies are
+        ["s3-fifo", "sieve", "lru", "clock"]. Default is "sieve".
+    pin_memory : bool, optional
+        Whether the cache storage should be allocated on system pinned memory.
+        Default is False.
+    Returns
+    -------
+    Union[CPUCachedFeature, Dict[FeatureKey, CPUCachedFeature]]
+        New feature(s) wrapped with CPUCachedFeature.
+    """
+    return wrap_with_cached_feature(
+        CPUCachedFeature,
+        fallback_features,
+        max_cache_size_in_bytes,
+        policy=policy,
+        pin_memory=pin_memory,
+    )

python/dgl/graphbolt/impl/cpu_feature_cache.py

@@ -0,0 +1,176 @@
+"""CPU Feature Cache implementation wrapper for graphbolt."""
+import torch
+
+__all__ = ["CPUFeatureCache"]
+
+caching_policies = {
+    "s3-fifo": torch.ops.graphbolt.s3_fifo_cache_policy,
+    "sieve": torch.ops.graphbolt.sieve_cache_policy,
+    "lru": torch.ops.graphbolt.lru_cache_policy,
+    "clock": torch.ops.graphbolt.clock_cache_policy,
+}
+
+
+class CPUFeatureCache(object):
+    r"""High level wrapper for the CPU feature cache.
+
+    Parameters
+    ----------
+    cache_shape : List[int]
+        The shape of the cache. cache_shape[0] gives us the capacity.
+    dtype : torch.dtype
+        The data type of the elements stored in the cache.
+    policy: str, optional
+        The cache policy. Default is "sieve". "s3-fifo", "lru" and "clock" are
+        also available.
+    num_parts: int, optional
+        The number of cache partitions for parallelism. Default is
+        `torch.get_num_threads()`.
+    pin_memory: bool, optional
+        Whether the cache storage should be pinned.
+    """
+
+    def __init__(
+        self,
+        cache_shape,
+        dtype,
+        policy=None,
+        num_parts=None,
+        pin_memory=False,
+    ):
+        if policy is None:
+            policy = "sieve"
+        assert (
+            policy in caching_policies
+        ), f"{list(caching_policies.keys())} are the available caching policies."
+        if num_parts is None:
+            num_parts = torch.get_num_threads()
+        min_num_cache_items = num_parts * (10 if policy == "s3-fifo" else 1)
+        # Since we partition the cache, each partition needs to have a positive
+        # number of slots. In addition, each "s3-fifo" partition needs at least
+        # 10 slots since the small queue is 10% and the small queue needs a
+        # positive size.
+        if cache_shape[0] < min_num_cache_items:
+            cache_shape = (min_num_cache_items,) + cache_shape[1:]
+        self._policy = caching_policies[policy](cache_shape[0], num_parts)
+        self._cache = torch.ops.graphbolt.feature_cache(
+            cache_shape, dtype, pin_memory
+        )
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def is_pinned(self):
+        """Returns True if the cache storage is pinned."""
+        return self._cache.is_pinned()
+
+    @property
+    def max_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._cache.nbytes
+
+    def query(self, keys, offset=0):
+        """Queries the cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the cache with.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+
+        Returns
+        -------
+        tuple(Tensor, Tensor, Tensor, Tensor)
+            A tuple containing
+            (values, missing_indices, missing_keys, missing_offsets) where
+            values[missing_indices] corresponds to cache misses that should be
+            filled by quering another source with missing_keys. If keys is
+            pinned, then the returned values tensor is pinned as well. The
+            missing_offsets tensor has the partition offsets of missing_keys.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            positions,
+            index,
+            missing_keys,
+            found_pointers,
+            found_offsets,
+            missing_offsets,
+        ) = self._policy.query(keys, offset)
+        values = self._cache.query(positions, index, keys.shape[0])
+        self._policy.reading_completed(found_pointers, found_offsets)
+        self.total_miss += missing_keys.shape[0]
+        missing_index = index[positions.size(0) :]
+        return values, missing_index, missing_keys, missing_offsets
+
+    def query_and_replace(self, keys, reader_fn, offset=0):
+        """Queries the cache. Then inserts the keys that are not found by
+        reading them by calling `reader_fn(missing_keys)`, which are then
+        inserted into the cache using the selected caching policy algorithm
+        to remove the old entries if it is full.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the cache with.
+        reader_fn : reader_fn(keys: torch.Tensor) -> torch.Tensor
+            A function that will take a missing keys tensor and will return
+            their values.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+
+        Returns
+        -------
+        Tensor
+            A tensor containing values corresponding to the keys. Should equal
+            `reader_fn(keys)`, computed in a faster way.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            positions,
+            index,
+            pointers,
+            missing_keys,
+            found_offsets,
+            missing_offsets,
+        ) = self._policy.query_and_replace(keys, offset)
+        found_cnt = keys.size(0) - missing_keys.size(0)
+        found_positions = positions[:found_cnt]
+        values = self._cache.query(found_positions, index, keys.shape[0])
+        found_pointers = pointers[:found_cnt]
+        self._policy.reading_completed(found_pointers, found_offsets)
+        self.total_miss += missing_keys.shape[0]
+        missing_index = index[found_cnt:]
+        missing_values = reader_fn(missing_keys)
+        values[missing_index] = missing_values
+        missing_positions = positions[found_cnt:]
+        self._cache.replace(missing_positions, missing_values)
+        missing_pointers = pointers[found_cnt:]
+        self._policy.writing_completed(missing_pointers, missing_offsets)
+        return values
+
+    def replace(self, keys, values, offsets=None, offset=0):
+        """Inserts key-value pairs into the cache using the selected caching
+        policy algorithm to remove old key-value pairs if it is full.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to insert to the cache.
+        values : Tensor
+            The values to insert to the cache.
+        offsets : Tensor, optional
+            The partition offsets of the keys.
+        offset : int
+            The offset to be added to the keys. Default is 0.
+        """
+        positions, pointers, offsets = self._policy.replace(
+            keys, offsets, offset
+        )
+        self._cache.replace(positions, values)
+        self._policy.writing_completed(pointers, offsets)
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/gpu_cached_feature.py

@@ -0,0 +1,283 @@
+"""GPU cached feature for GraphBolt."""
+from typing import Dict, Union
+
+import torch
+
+from ..feature_store import (
+    bytes_to_number_of_items,
+    Feature,
+    FeatureKey,
+    wrap_with_cached_feature,
+)
+
+from .gpu_feature_cache import GPUFeatureCache
+
+__all__ = ["GPUCachedFeature", "gpu_cached_feature"]
+
+
+class GPUCachedFeature(Feature):
+    r"""GPU cached feature wrapping a fallback feature. It uses the least
+    recently used (LRU) algorithm as the cache eviction policy. Use
+    `gpu_cached_feature` to construct an instance of this class.
+
+    Places the GPU cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    fallback_feature : Feature
+        The fallback feature.
+    cache : GPUFeatureCache
+        A GPUFeatureCache instance to serve as the cache backend.
+    offset : int, optional
+        The offset value to add to the given ids before using the cache. This
+        parameter is useful if multiple `GPUCachedFeature`s are sharing a single
+        GPUFeatureCache object.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> torch_feat = torch.arange(10).reshape(2, -1).to("cuda")
+    >>> cache_size = 5
+    >>> fallback_feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature = gb.gpu_cached_feature(fallback_feature, cache_size)
+    >>> feature.read()
+    tensor([[0, 1, 2, 3, 4],
+            [5, 6, 7, 8, 9]], device='cuda:0')
+    >>> feature.read(torch.tensor([0]).to("cuda"))
+    tensor([[0, 1, 2, 3, 4]], device='cuda:0')
+    >>> feature.update(torch.tensor([[1 for _ in range(5)]]).to("cuda"),
+    ...                torch.tensor([1]).to("cuda"))
+    >>> feature.read(torch.tensor([0, 1]).to("cuda"))
+    tensor([[0, 1, 2, 3, 4],
+            [1, 1, 1, 1, 1]], device='cuda:0')
+    >>> feature.size()
+    torch.Size([5])
+    """
+
+    _cache_type = GPUFeatureCache
+
+    def __init__(
+        self,
+        fallback_feature: Feature,
+        cache: GPUFeatureCache,
+        offset: int = 0,
+    ):
+        super(GPUCachedFeature, self).__init__()
+        assert isinstance(fallback_feature, Feature), (
+            f"The fallback_feature must be an instance of Feature, but got "
+            f"{type(fallback_feature)}."
+        )
+        self._fallback_feature = fallback_feature
+        self._feature = cache
+        self._offset = offset
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        The returned tensor is always in GPU memory, no matter whether the
+        fallback feature is in memory or on disk.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._fallback_feature.read()
+        values, missing_index, missing_keys = self._feature.query(
+            ids if self._offset == 0 else ids + self._offset
+        )
+        missing_values = self._fallback_feature.read(
+            missing_keys if self._offset == 0 else missing_keys - self._offset
+        )
+        values[missing_index] = missing_values
+        self._feature.replace(missing_keys, missing_values)
+        return values
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        future = self._feature.query(
+            ids if self._offset == 0 else ids + self._offset, async_op=True
+        )
+
+        yield
+
+        values, missing_index, missing_keys = future.wait()
+
+        fallback_reader = self._fallback_feature.read_async(
+            missing_keys if self._offset == 0 else missing_keys - self._offset
+        )
+        fallback_num_stages = self._fallback_feature.read_async_num_stages(
+            missing_keys.device
+        )
+        for i in range(fallback_num_stages):
+            missing_values_future = next(fallback_reader, None)
+            if i < fallback_num_stages - 1:
+                yield  # fallback feature stages.
+
+        class _Waiter:
+            def __init__(
+                self,
+                feature,
+                values,
+                missing_index,
+                missing_keys,
+                missing_values_future,
+            ):
+                self.feature = feature
+                self.values = values
+                self.missing_index = missing_index
+                self.missing_keys = missing_keys
+                self.missing_values_future = missing_values_future
+
+            def wait(self):
+                """Returns the stored value when invoked."""
+                missing_values = self.missing_values_future.wait()
+                self.feature.replace(self.missing_keys, missing_values)
+                self.values[self.missing_index] = missing_values
+                values = self.values
+                # Ensure there is no memory leak.
+                self.feature = self.values = self.missing_index = None
+                self.missing_keys = self.missing_values_future = None
+                return values
+
+        yield _Waiter(
+            self._feature,
+            values,
+            missing_index,
+            missing_keys,
+            missing_values_future,
+        )
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        assert ids_device.type == "cuda"
+        return 1 + self._fallback_feature.read_async_num_stages(ids_device)
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._fallback_feature.size()
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._fallback_feature.count()
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            feat0 = value[:1]
+            self._fallback_feature.update(value)
+            cache_size = min(
+                bytes_to_number_of_items(self.cache_size_in_bytes, feat0),
+                value.shape[0],
+            )
+            self._feature = None  # Destroy the existing cache first.
+            self._feature = self._cache_type(
+                (cache_size,) + feat0.shape[1:], feat0.dtype
+            )
+        else:
+            self._fallback_feature.update(value, ids)
+            self._feature.replace(ids, value)
+
+    @property
+    def cache_size_in_bytes(self):
+        """Return the size taken by the cache in bytes."""
+        return self._feature.max_size_in_bytes
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self._feature.miss_rate
+
+
+def gpu_cached_feature(
+    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
+    max_cache_size_in_bytes: int,
+) -> Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]:
+    r"""GPU cached feature wrapping a fallback feature. It uses the least
+    recently used (LRU) algorithm as the cache eviction policy.
+
+    Places the GPU cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
+        The fallback feature(s).
+    max_cache_size_in_bytes : int
+        The capacity of the GPU cache in bytes.
+    Returns
+    -------
+    Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]
+        The feature(s) wrapped with GPUCachedFeature.
+    """
+    return wrap_with_cached_feature(
+        GPUCachedFeature, fallback_features, max_cache_size_in_bytes
+    )

python/dgl/graphbolt/impl/gpu_feature_cache.py

@@ -0,0 +1,80 @@
+"""HugeCTR gpu_cache wrapper for graphbolt."""
+from functools import reduce
+from operator import mul
+
+import torch
+
+
+class GPUFeatureCache(object):
+    """High-level wrapper for GPU embedding cache"""
+
+    def __init__(self, cache_shape, dtype):
+        major, _ = torch.cuda.get_device_capability()
+        assert (
+            major >= 7
+        ), "GPUFeatureCache is supported only on CUDA compute capability >= 70 (Volta)."
+        self._cache = torch.ops.graphbolt.gpu_cache(cache_shape, dtype)
+        element_size = torch.tensor([], dtype=dtype).element_size()
+        self.max_size_in_bytes = reduce(mul, cache_shape) * element_size
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def query(self, keys, async_op=False):
+        """Queries the GPU cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU cache with.
+        async_op: bool
+            Boolean indicating whether the call is asynchronous. If so, the
+            result can be obtained by calling wait on the returned future.
+
+        Returns
+        -------
+        tuple(Tensor, Tensor, Tensor)
+            A tuple containing (values, missing_indices, missing_keys) where
+            values[missing_indices] corresponds to cache misses that should be
+            filled by quering another source with missing_keys.
+        """
+
+        class _Waiter:
+            def __init__(self, gpu_cache, future):
+                self.gpu_cache = gpu_cache
+                self.future = future
+
+            def wait(self):
+                """Returns the stored value when invoked."""
+                gpu_cache = self.gpu_cache
+                values, missing_index, missing_keys = (
+                    self.future.wait() if async_op else self.future
+                )
+                # Ensure there is no leak.
+                self.gpu_cache = self.future = None
+
+                gpu_cache.total_queries += values.shape[0]
+                gpu_cache.total_miss += missing_keys.shape[0]
+                return values, missing_index, missing_keys
+
+        if async_op:
+            return _Waiter(self, self._cache.query_async(keys))
+        else:
+            return _Waiter(self, self._cache.query(keys)).wait()
+
+    def replace(self, keys, values):
+        """Inserts key-value pairs into the GPU cache using the Least-Recently
+        Used (LRU) algorithm to remove old key-value pairs if it is full.
+
+        Parameters
+        ----------
+        keys: Tensor
+            The keys to insert to the GPU cache.
+        values: Tensor
+            The values to insert to the GPU cache.
+        """
+        self._cache.replace(keys, values)
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/gpu_graph_cache.py

@@ -0,0 +1,118 @@
+"""HugeCTR gpu_cache wrapper for graphbolt."""
+import torch
+
+
+class GPUGraphCache(object):
+    r"""High-level wrapper for GPU graph cache.
+
+    Places the GPU graph cache to torch.cuda.current_device().
+
+    Parameters
+    ----------
+    num_edges : int
+        Upperbound on number of edges to cache.
+    threshold : int
+        The number of accesses before the neighborhood of a vertex is cached.
+    indptr_dtype : torch.dtype
+        The dtype of the indptr tensor of the graph.
+    dtypes : list[torch.dtype]
+        The dtypes of the edge tensors that are going to be cached.
+    has_original_edge_ids : bool
+        Whether the graph to be cached has original edge ids.
+    """
+
+    def __init__(
+        self, num_edges, threshold, indptr_dtype, dtypes, has_original_edge_ids
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        assert (
+            major >= 7
+        ), "GPUGraphCache is supported only on CUDA compute capability >= 70 (Volta)."
+        self._cache = torch.ops.graphbolt.gpu_graph_cache(
+            num_edges, threshold, indptr_dtype, dtypes, has_original_edge_ids
+        )
+        self.total_miss = 0
+        self.total_queries = 0
+
+    def query(self, keys):
+        """Queries the GPU cache.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU graph cache with.
+
+        Returns
+        -------
+        tuple(Tensor, func)
+            A tuple containing (missing_keys, replace_fn) where replace_fn is a
+            function that should be called with the graph structure
+            corresponding to the missing keys. Its arguments are
+            (Tensor, list(Tensor)), where the first tensor is the missing indptr
+            and the second list is the missing edge tensors.
+        """
+        self.total_queries += keys.shape[0]
+        (
+            index,
+            position,
+            num_hit,
+            num_threshold,
+        ) = self._cache.query(keys)
+        self.total_miss += keys.shape[0] - num_hit
+
+        def replace_functional(missing_indptr, missing_edge_tensors):
+            return self._cache.replace(
+                keys,
+                index,
+                position,
+                num_hit,
+                num_threshold,
+                missing_indptr,
+                missing_edge_tensors,
+            )
+
+        return keys[index[num_hit:]], replace_functional
+
+    def query_async(self, keys):
+        """Queries the GPU cache asynchronously.
+
+        Parameters
+        ----------
+        keys : Tensor
+            The keys to query the GPU graph cache with.
+
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns the missing keys on the second
+            invocation and expects the fetched indptr and edge tensors on the
+            next invocation. The third and last invocation returns a future
+            object and the return result can be accessed by calling `.wait()`
+            on the returned future object. It is undefined behavior to call
+            `.wait()` more than once.
+        """
+        future = self._cache.query_async(keys)
+
+        yield
+
+        index, position, num_hit, num_threshold = future.wait()
+
+        self.total_queries += keys.shape[0]
+        self.total_miss += keys.shape[0] - num_hit
+
+        missing_indptr, missing_edge_tensors = yield keys[index[num_hit:]]
+
+        yield self._cache.replace_async(
+            keys,
+            index,
+            position,
+            num_hit,
+            num_threshold,
+            missing_indptr,
+            missing_edge_tensors,
+        )
+
+    @property
+    def miss_rate(self):
+        """Returns the cache miss rate since creation."""
+        return self.total_miss / self.total_queries

python/dgl/graphbolt/impl/in_subgraph_sampler.py

@@ -0,0 +1,86 @@
+"""In-subgraph sampler for GraphBolt."""
+
+from torch.utils.data import functional_datapipe
+
+from ..internal import unique_and_compact_csc_formats
+
+from ..subgraph_sampler import SubgraphSampler
+from .sampled_subgraph_impl import SampledSubgraphImpl
+
+
+__all__ = ["InSubgraphSampler"]
+
+
+@functional_datapipe("sample_in_subgraph")
+class InSubgraphSampler(SubgraphSampler):
+    """Sample the subgraph induced on the inbound edges of the given nodes.
+
+    Functional name: :obj:`sample_in_subgraph`.
+
+    In-subgraph sampler is responsible for sampling a subgraph from given data,
+    returning an induced subgraph along with compacted information.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform in_subgraph sampling.
+
+    Examples
+    -------
+    >>> import dgl.graphbolt as gb
+    >>> import torch
+    >>> indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
+    >>> indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
+    >>> graph = gb.fused_csc_sampling_graph(indptr, indices)
+    >>> item_set = gb.ItemSet(len(indptr) - 1, names="seeds")
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=2)
+    >>> insubgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
+    >>> for _, data in enumerate(insubgraph_sampler):
+    ...     print(data.sampled_subgraphs[0].sampled_csc)
+    ...     print(data.sampled_subgraphs[0].original_row_node_ids)
+    ...     print(data.sampled_subgraphs[0].original_column_node_ids)
+    CSCFormatBase(indptr=tensor([0, 3, 5]),
+                indices=tensor([0, 1, 2, 3, 4]),
+    )
+    tensor([0, 1, 4, 2, 3])
+    tensor([0, 1])
+    CSCFormatBase(indptr=tensor([0, 2, 4]),
+                indices=tensor([2, 3, 4, 0]),
+    )
+    tensor([2, 3, 0, 5, 1])
+    tensor([2, 3])
+    CSCFormatBase(indptr=tensor([0, 3, 5]),
+                indices=tensor([2, 3, 1, 4, 0]),
+    )
+    tensor([4, 5, 0, 3, 1])
+    tensor([4, 5])
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+    ):
+        super().__init__(datapipe)
+        self.graph = graph
+        self.sampler = graph.in_subgraph
+
+    def sample_subgraphs(
+        self, seeds, seeds_timestamp, seeds_pre_time_window=None
+    ):
+        subgraph = self.sampler(seeds)
+        (
+            original_row_node_ids,
+            compacted_csc_formats,
+            _,
+        ) = unique_and_compact_csc_formats(subgraph.sampled_csc, seeds)
+        subgraph = SampledSubgraphImpl(
+            sampled_csc=compacted_csc_formats,
+            original_column_node_ids=seeds,
+            original_row_node_ids=original_row_node_ids,
+            original_edge_ids=subgraph.original_edge_ids,
+        )
+        seeds = original_row_node_ids
+        return (seeds, [subgraph])

python/dgl/graphbolt/impl/legacy_dataset.py

@@ -0,0 +1,158 @@
+"""Graphbolt dataset for legacy DGLDataset."""
+
+from typing import List, Union
+
+from ..base import etype_tuple_to_str
+from ..dataset import Dataset, Task
+from ..itemset import HeteroItemSet, ItemSet
+from ..sampling_graph import SamplingGraph
+from .basic_feature_store import BasicFeatureStore
+from .fused_csc_sampling_graph import from_dglgraph
+from .ondisk_dataset import OnDiskTask
+from .torch_based_feature_store import TorchBasedFeature
+
+
+class LegacyDataset(Dataset):
+    """A Graphbolt dataset for legacy DGLDataset."""
+
+    def __init__(self, legacy):
+        # Only supports single graph cases.
+        assert len(legacy) == 1
+        graph = legacy[0]
+        # Handle OGB Dataset.
+        if isinstance(graph, tuple):
+            graph, _ = graph
+        if graph.is_homogeneous:
+            self._init_as_homogeneous_node_pred(legacy)
+        else:
+            self._init_as_heterogeneous_node_pred(legacy)
+
+    def _init_as_heterogeneous_node_pred(self, legacy):
+        def _init_item_set_dict(idx, labels):
+            item_set_dict = {}
+            for key in idx.keys():
+                item_set = ItemSet(
+                    (idx[key], labels[key][idx[key]]),
+                    names=("seeds", "labels"),
+                )
+                item_set_dict[key] = item_set
+            return HeteroItemSet(item_set_dict)
+
+        # OGB Dataset has the idx split.
+        if hasattr(legacy, "get_idx_split"):
+            graph, labels = legacy[0]
+            split_idx = legacy.get_idx_split()
+
+            # Initialize tasks.
+            tasks = []
+            metadata = {
+                "num_classes": legacy.num_classes,
+                "name": "node_classification",
+            }
+            train_set = _init_item_set_dict(split_idx["train"], labels)
+            validation_set = _init_item_set_dict(split_idx["valid"], labels)
+            test_set = _init_item_set_dict(split_idx["test"], labels)
+            task = OnDiskTask(metadata, train_set, validation_set, test_set)
+            tasks.append(task)
+            self._tasks = tasks
+
+            item_set_dict = {}
+            for ntype in graph.ntypes:
+                item_set = ItemSet(graph.num_nodes(ntype), names="seeds")
+                item_set_dict[ntype] = item_set
+            self._all_nodes_set = HeteroItemSet(item_set_dict)
+
+            features = {}
+            for ntype in graph.ntypes:
+                for name in graph.nodes[ntype].data.keys():
+                    tensor = graph.nodes[ntype].data[name]
+                    if tensor.dim() == 1:
+                        tensor = tensor.view(-1, 1)
+                    features[("node", ntype, name)] = TorchBasedFeature(tensor)
+            for etype in graph.canonical_etypes:
+                for name in graph.edges[etype].data.keys():
+                    tensor = graph.edges[etype].data[name]
+                    if tensor.dim() == 1:
+                        tensor = tensor.view(-1, 1)
+                    gb_etype = etype_tuple_to_str(etype)
+                    features[("edge", gb_etype, name)] = TorchBasedFeature(
+                        tensor
+                    )
+            self._feature = BasicFeatureStore(features)
+            self._graph = from_dglgraph(graph, is_homogeneous=False)
+            self._dataset_name = legacy.name
+        else:
+            raise NotImplementedError(
+                "Only support heterogeneous ogn node pred dataset"
+            )
+
+    def _init_as_homogeneous_node_pred(self, legacy):
+        from dgl.data import AsNodePredDataset
+
+        legacy = AsNodePredDataset(legacy)
+
+        # Initialize tasks.
+        tasks = []
+        metadata = {
+            "num_classes": legacy.num_classes,
+            "name": "node_classification",
+        }
+        train_labels = legacy[0].ndata["label"][legacy.train_idx]
+        validation_labels = legacy[0].ndata["label"][legacy.val_idx]
+        test_labels = legacy[0].ndata["label"][legacy.test_idx]
+        train_set = ItemSet(
+            (legacy.train_idx, train_labels),
+            names=("seeds", "labels"),
+        )
+        validation_set = ItemSet(
+            (legacy.val_idx, validation_labels),
+            names=("seeds", "labels"),
+        )
+        test_set = ItemSet(
+            (legacy.test_idx, test_labels), names=("seeds", "labels")
+        )
+        task = OnDiskTask(metadata, train_set, validation_set, test_set)
+        tasks.append(task)
+        self._tasks = tasks
+
+        num_nodes = legacy[0].num_nodes()
+        self._all_nodes_set = ItemSet(num_nodes, names="seeds")
+        features = {}
+        for name in legacy[0].ndata.keys():
+            tensor = legacy[0].ndata[name]
+            if tensor.dim() == 1:
+                tensor = tensor.view(-1, 1)
+            features[("node", None, name)] = TorchBasedFeature(tensor)
+        for name in legacy[0].edata.keys():
+            tensor = legacy[0].edata[name]
+            if tensor.dim() == 1:
+                tensor = tensor.view(-1, 1)
+            features[("edge", None, name)] = TorchBasedFeature(tensor)
+        self._feature = BasicFeatureStore(features)
+        self._graph = from_dglgraph(legacy[0], is_homogeneous=True)
+        self._dataset_name = legacy.name
+
+    @property
+    def tasks(self) -> List[Task]:
+        """Return the tasks."""
+        return self._tasks
+
+    @property
+    def graph(self) -> SamplingGraph:
+        """Return the graph."""
+        return self._graph
+
+    @property
+    def feature(self) -> BasicFeatureStore:
+        """Return the feature."""
+        return self._feature
+
+    @property
+    def dataset_name(self) -> str:
+        """Return the dataset name."""
+        return self._dataset_name
+
+    @property
+    def all_nodes_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the itemset containing all nodes."""
+        return self._all_nodes_set

python/dgl/graphbolt/impl/ondisk_metadata.py

@@ -0,0 +1,129 @@
+"""Ondisk metadata of GraphBolt."""
+
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import pydantic
+
+from ..internal_utils import version
+
+
+__all__ = [
+    "OnDiskFeatureDataFormat",
+    "OnDiskTVTSetData",
+    "OnDiskTVTSet",
+    "OnDiskFeatureDataDomain",
+    "OnDiskFeatureData",
+    "OnDiskMetaData",
+    "OnDiskGraphTopologyType",
+    "OnDiskGraphTopology",
+    "OnDiskTaskData",
+]
+
+
+class ExtraMetaData(pydantic.BaseModel, extra="allow"):
+    """Group extra fields into metadata. Internal use only."""
+
+    extra_fields: Optional[Dict[str, Any]] = {}
+
+    # As pydantic 2.0 has changed the API of validators, we need to use
+    # different validators for different versions to be compatible with
+    # previous versions.
+    if version.parse(pydantic.__version__) >= version.parse("2.0"):
+
+        @pydantic.model_validator(mode="before")
+        @classmethod
+        def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+            """Build extra fields."""
+            for key in list(values.keys()):
+                if key not in cls.model_fields:
+                    values["extra_fields"] = values.get("extra_fields", {})
+                    values["extra_fields"][key] = values.pop(key)
+            return values
+
+    else:
+
+        @pydantic.root_validator(pre=True)
+        @classmethod
+        def build_extra(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+            """Build extra fields."""
+            for key in list(values.keys()):
+                if key not in cls.__fields__:
+                    values["extra_fields"] = values.get("extra_fields", {})
+                    values["extra_fields"][key] = values.pop(key)
+            return values
+
+
+class OnDiskFeatureDataFormat(str, Enum):
+    """Enum of data format."""
+
+    TORCH = "torch"
+    NUMPY = "numpy"
+
+
+class OnDiskTVTSetData(pydantic.BaseModel):
+    """Train-Validation-Test set data."""
+
+    name: Optional[str] = None
+    format: OnDiskFeatureDataFormat
+    in_memory: Optional[bool] = True
+    path: str
+
+
+class OnDiskTVTSet(pydantic.BaseModel):
+    """Train-Validation-Test set."""
+
+    type: Optional[str] = None
+    data: List[OnDiskTVTSetData]
+
+
+class OnDiskFeatureDataDomain(str, Enum):
+    """Enum of feature data domain."""
+
+    NODE = "node"
+    EDGE = "edge"
+    GRAPH = "graph"
+
+
+class OnDiskFeatureData(ExtraMetaData):
+    r"""The description of an on-disk feature."""
+    domain: OnDiskFeatureDataDomain
+    type: Optional[str] = None
+    name: str
+    format: OnDiskFeatureDataFormat
+    path: str
+    in_memory: Optional[bool] = True
+
+
+class OnDiskGraphTopologyType(str, Enum):
+    """Enum of graph topology type."""
+
+    FUSED_CSC_SAMPLING = "FusedCSCSamplingGraph"
+
+
+class OnDiskGraphTopology(pydantic.BaseModel):
+    """The description of an on-disk graph topology."""
+
+    type: OnDiskGraphTopologyType
+    path: str
+
+
+class OnDiskTaskData(ExtraMetaData):
+    """Task specification in YAML."""
+
+    train_set: Optional[List[OnDiskTVTSet]] = []
+    validation_set: Optional[List[OnDiskTVTSet]] = []
+    test_set: Optional[List[OnDiskTVTSet]] = []
+
+
+class OnDiskMetaData(pydantic.BaseModel):
+    """Metadata specification in YAML.
+
+    As multiple node/edge types and multiple splits are supported, each TVT set
+    is a list of list of ``OnDiskTVTSet``.
+    """
+
+    dataset_name: Optional[str] = None
+    graph_topology: Optional[OnDiskGraphTopology] = None
+    feature_data: Optional[List[OnDiskFeatureData]] = []
+    tasks: Optional[List[OnDiskTaskData]] = []

python/dgl/graphbolt/impl/sampled_subgraph_impl.py

@@ -0,0 +1,118 @@
+"""Sampled subgraph for FusedCSCSamplingGraph."""
+# pylint: disable= invalid-name
+from dataclasses import dataclass
+from typing import Dict, Union
+
+import torch
+
+from ..base import CSCFormatBase, etype_str_to_tuple
+from ..internal_utils import get_attributes
+from ..sampled_subgraph import SampledSubgraph
+
+__all__ = ["SampledSubgraphImpl"]
+
+
+@dataclass
+class SampledSubgraphImpl(SampledSubgraph):
+    r"""Sampled subgraph of CSCSamplingGraph.
+
+    Examples
+    --------
+    >>> sampled_csc = {"A:relation:B": CSCFormatBase(indptr=torch.tensor([0, 1, 2, 3]),
+    ... indices=torch.tensor([0, 1, 2]))}
+    >>> original_column_node_ids = {'B': torch.tensor([10, 11, 12])}
+    >>> original_row_node_ids = {'A': torch.tensor([13, 14, 15])}
+    >>> original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+    >>> subgraph = gb.SampledSubgraphImpl(
+    ... sampled_csc=sampled_csc,
+    ... original_column_node_ids=original_column_node_ids,
+    ... original_row_node_ids=original_row_node_ids,
+    ... original_edge_ids=original_edge_ids
+    ... )
+    >>> print(subgraph.sampled_csc)
+    {"A:relation:B": CSCForamtBase(indptr=torch.tensor([0, 1, 2, 3]),
+    ... indices=torch.tensor([0, 1, 2]))}
+    >>> print(subgraph.original_column_node_ids)
+    {'B': tensor([10, 11, 12])}
+    >>> print(subgraph.original_row_node_ids)
+    {'A': tensor([13, 14, 15])}
+    >>> print(subgraph.original_edge_ids)
+    {"A:relation:B": tensor([19, 20, 21])}
+    """
+    sampled_csc: Union[CSCFormatBase, Dict[str, CSCFormatBase]] = None
+    original_column_node_ids: Union[
+        Dict[str, torch.Tensor], torch.Tensor
+    ] = None
+    original_row_node_ids: Union[Dict[str, torch.Tensor], torch.Tensor] = None
+    original_edge_ids: Union[Dict[str, torch.Tensor], torch.Tensor] = None
+    # Used to fetch sampled_csc.indices if it is missing.
+    _edge_ids_in_fused_csc_sampling_graph: Union[
+        Dict[str, torch.Tensor], torch.Tensor
+    ] = None
+
+    def __post_init__(self):
+        if isinstance(self.sampled_csc, dict):
+            for etype, pair in self.sampled_csc.items():
+                assert (
+                    isinstance(etype, str)
+                    and len(etype_str_to_tuple(etype)) == 3
+                ), "Edge type should be a string in format of str:str:str."
+                assert pair.indptr is not None and isinstance(
+                    pair.indptr, torch.Tensor
+                ), "Node pair should be have indptr of type torch.Tensor."
+                # For CUDA, indices may be None because it will be fetched later.
+                if not pair.indptr.is_cuda or pair.indices is not None:
+                    assert isinstance(
+                        pair.indices, torch.Tensor
+                    ), "Node pair should be have indices of type torch.Tensor."
+                else:
+                    assert isinstance(
+                        self._edge_ids_in_fused_csc_sampling_graph.get(
+                            etype, None
+                        ),
+                        torch.Tensor,
+                    ), "When indices is missing, sampled edge ids needs to be provided."
+        else:
+            assert self.sampled_csc.indptr is not None and isinstance(
+                self.sampled_csc.indptr, torch.Tensor
+            ), "Node pair should be have torch.Tensor indptr."
+            # For CUDA, indices may be None because it will be fetched later.
+            if (
+                not self.sampled_csc.indptr.is_cuda
+                or self.sampled_csc.indices is not None
+            ):
+                assert isinstance(
+                    self.sampled_csc.indices, torch.Tensor
+                ), "Node pair should have a torch.Tensor indices."
+            else:
+                assert isinstance(
+                    self._edge_ids_in_fused_csc_sampling_graph, torch.Tensor
+                ), "When indices is missing, sampled edge ids needs to be provided."
+
+    def __repr__(self) -> str:
+        return _sampled_subgraph_str(self, "SampledSubgraphImpl")
+
+
+def _sampled_subgraph_str(sampled_subgraph: SampledSubgraph, classname) -> str:
+    final_str = classname + "("
+
+    attributes = get_attributes(sampled_subgraph)
+    attributes.reverse()
+
+    for name in attributes:
+        if name in "_edge_ids_in_fused_csc_sampling_graph":
+            continue
+        val = getattr(sampled_subgraph, name)
+
+        def _add_indent(_str, indent):
+            lines = _str.split("\n")
+            lines = [lines[0]] + [" " * indent + line for line in lines[1:]]
+            return "\n".join(lines)
+
+        val = str(val)
+        final_str = (
+            final_str
+            + f"{name}={_add_indent(val, len(name) + len(classname) + 1)},\n"
+            + " " * len(classname)
+        )
+    return final_str[: -len(classname)] + ")"

python/dgl/graphbolt/impl/temporal_neighbor_sampler.py

@@ -0,0 +1,270 @@
+"""Temporal neighbor subgraph samplers for GraphBolt."""
+import torch
+from torch.utils.data import functional_datapipe
+
+from ..internal import compact_csc_format
+
+from ..subgraph_sampler import SubgraphSampler
+from .sampled_subgraph_impl import SampledSubgraphImpl
+
+
+__all__ = ["TemporalNeighborSampler", "TemporalLayerNeighborSampler"]
+
+
+class TemporalNeighborSamplerImpl(SubgraphSampler):
+    """Base class for TemporalNeighborSamplers."""
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace,
+        prob_name,
+        node_timestamp_attr_name,
+        edge_timestamp_attr_name,
+        sampler,
+    ):
+        super().__init__(datapipe)
+        self.graph = graph
+        # Convert fanouts to a list of tensors.
+        self.fanouts = []
+        for fanout in fanouts:
+            if not isinstance(fanout, torch.Tensor):
+                fanout = torch.LongTensor([int(fanout)])
+            self.fanouts.insert(0, fanout)
+        self.replace = replace
+        self.prob_name = prob_name
+        self.node_timestamp_attr_name = node_timestamp_attr_name
+        self.edge_timestamp_attr_name = edge_timestamp_attr_name
+        self.sampler = sampler
+
+    def sample_subgraphs(
+        self, seeds, seeds_timestamp, seeds_pre_time_window=None
+    ):
+        assert (
+            seeds_timestamp is not None
+        ), "seeds_timestamp must be provided for temporal neighbor sampling."
+        subgraphs = []
+        num_layers = len(self.fanouts)
+        # Enrich seeds with all node types. Ensure that the dtype and device
+        # remain consistent with those of the existing seeds.
+        if isinstance(seeds, dict):
+            first_val = next(iter(seeds.items()))[1]
+            ntypes = list(self.graph.node_type_to_id.keys())
+            seeds = {
+                ntype: seeds.get(
+                    ntype,
+                    torch.tensor(
+                        [], dtype=first_val.dtype, device=first_val.device
+                    ),
+                )
+                for ntype in ntypes
+            }
+            empty_tensor = torch.tensor(
+                [], dtype=torch.int64, device=first_val.device
+            )
+            seeds_timestamp = {
+                ntype: seeds_timestamp.get(ntype, empty_tensor)
+                for ntype in ntypes
+            }
+            if seeds_pre_time_window:
+                seeds_pre_time_window = {
+                    ntype: seeds_pre_time_window.get(ntype, empty_tensor)
+                    for ntype in ntypes
+                }
+        for hop in range(num_layers):
+            subgraph = self.sampler(
+                seeds,
+                seeds_timestamp,
+                self.fanouts[hop],
+                self.replace,
+                seeds_pre_time_window,
+                self.prob_name,
+                self.node_timestamp_attr_name,
+                self.edge_timestamp_attr_name,
+            )
+            (
+                original_row_node_ids,
+                compacted_csc_formats,
+                row_timestamps,
+            ) = compact_csc_format(subgraph.sampled_csc, seeds, seeds_timestamp)
+
+            subgraph = SampledSubgraphImpl(
+                sampled_csc=compacted_csc_formats,
+                original_column_node_ids=seeds,
+                original_row_node_ids=original_row_node_ids,
+                original_edge_ids=subgraph.original_edge_ids,
+            )
+
+            subgraphs.insert(0, subgraph)
+            seeds = original_row_node_ids
+            seeds_timestamp = row_timestamps
+        return seeds, subgraphs
+
+
+@functional_datapipe("temporal_sample_neighbor")
+class TemporalNeighborSampler(TemporalNeighborSamplerImpl):
+    """Temporally sample neighbor edges from a graph and return sampled
+    subgraphs.
+
+    Functional name: :obj:`temporal_sample_neighbor`.
+
+    Neighbor sampler is responsible for sampling a subgraph from given data. It
+    returns an induced subgraph along with compacted information. In the
+    context of a node classification task, the neighbor sampler directly
+    utilizes the nodes provided as seed nodes. However, in scenarios involving
+    link prediction, the process needs another pre-peocess operation. That is,
+    gathering unique nodes from the given node pairs, encompassing both
+    positive and negative node pairs, and employs these nodes as the seed nodes
+    for subsequent steps.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform subgraph sampling.
+    fanouts: list[torch.Tensor] or list[int]
+        The number of edges to be sampled for each node with or without
+        considering edge types. The length of this parameter implicitly
+        signifies the layer of sampling being conducted.
+        Note: The fanout order is from the outermost layer to innermost layer.
+        For example, the fanout '[15, 10, 5]' means that 15 to the outermost
+        layer, 10 to the intermediate layer and 5 corresponds to the innermost
+        layer.
+    replace: bool
+        Boolean indicating whether the sample is preformed with or
+        without replacement. If True, a value can be selected multiple
+        times. Otherwise, each value can be selected only once.
+    prob_name: str, optional
+        The name of an edge attribute used as the weights of sampling for
+        each node. This attribute tensor should contain (unnormalized)
+        probabilities corresponding to each neighboring edge of a node.
+        It must be a 1D floating-point or boolean tensor, with the number
+        of elements equalling the total number of edges.
+    node_timestamp_attr_name: str, optional
+        The name of an node attribute used as the timestamps of nodes.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of nodes.
+    edge_timestamp_attr_name: str, optional
+        The name of an edge attribute used as the timestamps of edges.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of edges.
+
+    Examples
+    -------
+    TODO(zhenkun) : Add an example after the API to pass timestamps is finalized.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace=False,
+        prob_name=None,
+        node_timestamp_attr_name=None,
+        edge_timestamp_attr_name=None,
+    ):
+        super().__init__(
+            datapipe,
+            graph,
+            fanouts,
+            replace,
+            prob_name,
+            node_timestamp_attr_name,
+            edge_timestamp_attr_name,
+            graph.temporal_sample_neighbors,
+        )
+
+
+@functional_datapipe("temporal_sample_layer_neighbor")
+class TemporalLayerNeighborSampler(TemporalNeighborSamplerImpl):
+    """Temporally sample neighbor edges from a graph and return sampled
+    subgraphs.
+
+    Functional name: :obj:`temporal_sample_layer_neighbor`.
+
+    Sampler that builds computational dependency of node representations via
+    labor sampling for multilayer GNN from the NeurIPS 2023 paper
+    `Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs
+    <https://proceedings.neurips.cc/paper_files/paper/2023/file/51f9036d5e7ae822da8f6d4adda1fb39-Paper-Conference.pdf>`__
+
+    Layer-Neighbor sampler is responsible for sampling a subgraph from given
+    data. It returns an induced subgraph along with compacted information. In
+    the context of a node classification task, the neighbor sampler directly
+    utilizes the nodes provided as seed nodes. However, in scenarios involving
+    link prediction, the process needs another pre-process operation. That is,
+    gathering unique nodes from the given node pairs, encompassing both
+    positive and negative node pairs, and employs these nodes as the seed nodes
+    for subsequent steps. When the graph is hetero, sampled subgraphs in
+    minibatch will contain every edge type even though it is empty after
+    sampling.
+
+    Implements the approach described in Appendix A.3 of the paper. Similar to
+    dgl.dataloading.LaborSampler but this uses sequential poisson sampling
+    instead of poisson sampling to keep the count of sampled edges per vertex
+    deterministic like NeighborSampler. Thus, it is a drop-in replacement for
+    NeighborSampler. However, unlike NeighborSampler, it samples fewer vertices
+    and edges for multilayer GNN scenario without harming convergence speed with
+    respect to training iterations.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform subgraph sampling.
+    fanouts: list[torch.Tensor] or list[int]
+        The number of edges to be sampled for each node with or without
+        considering edge types. The length of this parameter implicitly
+        signifies the layer of sampling being conducted.
+        Note: The fanout order is from the outermost layer to innermost layer.
+        For example, the fanout '[15, 10, 5]' means that 15 to the outermost
+        layer, 10 to the intermediate layer and 5 corresponds to the innermost
+        layer.
+    replace: bool
+        Boolean indicating whether the sample is preformed with or
+        without replacement. If True, a value can be selected multiple
+        times. Otherwise, each value can be selected only once.
+    prob_name: str, optional
+        The name of an edge attribute used as the weights of sampling for
+        each node. This attribute tensor should contain (unnormalized)
+        probabilities corresponding to each neighboring edge of a node.
+        It must be a 1D floating-point or boolean tensor, with the number
+        of elements equalling the total number of edges.
+    node_timestamp_attr_name: str, optional
+        The name of an node attribute used as the timestamps of nodes.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of nodes.
+    edge_timestamp_attr_name: str, optional
+        The name of an edge attribute used as the timestamps of edges.
+        It must be a 1D integer tensor, with the number of elements
+        equalling the total number of edges.
+
+    Examples
+    -------
+    TODO(zhenkun) : Add an example after the API to pass timestamps is finalized.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        fanouts,
+        replace=False,
+        prob_name=None,
+        node_timestamp_attr_name=None,
+        edge_timestamp_attr_name=None,
+    ):
+        super().__init__(
+            datapipe,
+            graph,
+            fanouts,
+            replace,
+            prob_name,
+            node_timestamp_attr_name,
+            edge_timestamp_attr_name,
+            graph.temporal_sample_layer_neighbors,
+        )

python/dgl/graphbolt/impl/torch_based_feature_store.py

@@ -0,0 +1,665 @@
+"""Torch-based feature store for GraphBolt."""
+
+import copy
+import textwrap
+from typing import Dict, List
+
+import numpy as np
+import torch
+
+from ..base import (
+    get_device_to_host_uva_stream,
+    get_host_to_device_uva_stream,
+    index_select,
+)
+from ..feature_store import Feature
+from ..internal_utils import gb_warning, is_wsl
+from .basic_feature_store import BasicFeatureStore
+from .ondisk_metadata import OnDiskFeatureData
+
+__all__ = ["TorchBasedFeature", "DiskBasedFeature", "TorchBasedFeatureStore"]
+
+
+class _Waiter:
+    def __init__(self, event, values):
+        self.event = event
+        self.values = values
+
+    def wait(self):
+        """Returns the stored value when invoked."""
+        self.event.wait()
+        values = self.values
+        # Ensure there is no memory leak.
+        self.event = self.values = None
+        return values
+
+
+class TorchBasedFeature(Feature):
+    r"""A wrapper of pytorch based feature.
+
+    Initialize a torch based feature store by a torch feature.
+    Note that the feature can be either in memory or on disk.
+
+    Parameters
+    ----------
+    torch_feature : torch.Tensor
+        The torch feature.
+        Note that the dimension of the tensor should be greater than 1.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+
+    1. The feature is in memory.
+
+    >>> torch_feat = torch.arange(10).reshape(2, -1)
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read()
+    tensor([[0, 1, 2, 3, 4],
+            [5, 6, 7, 8, 9]])
+    >>> feature.read(torch.tensor([0]))
+    tensor([[0, 1, 2, 3, 4]])
+    >>> feature.update(torch.tensor([[1 for _ in range(5)]]),
+    ...                      torch.tensor([1]))
+    >>> feature.read(torch.tensor([0, 1]))
+    tensor([[0, 1, 2, 3, 4],
+            [1, 1, 1, 1, 1]])
+    >>> feature.size()
+    torch.Size([5])
+
+    2. The feature is on disk. Note that you can use gb.numpy_save_aligned as a
+    replacement for np.save to potentially get increased performance.
+
+    >>> import numpy as np
+    >>> arr = np.array([[1, 2], [3, 4]])
+    >>> np.save("/tmp/arr.npy", arr)
+    >>> torch_feat = torch.from_numpy(np.load("/tmp/arr.npy", mmap_mode="r+"))
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read()
+    tensor([[1, 2],
+            [3, 4]])
+    >>> feature.read(torch.tensor([0]))
+    tensor([[1, 2]])
+
+    3. Pinned CPU feature.
+
+    >>> torch_feat = torch.arange(10).reshape(2, -1).pin_memory()
+    >>> feature = gb.TorchBasedFeature(torch_feat)
+    >>> feature.read().device
+    device(type='cuda', index=0)
+    >>> feature.read(torch.tensor([0]).cuda()).device
+    device(type='cuda', index=0)
+    """
+
+    def __init__(self, torch_feature: torch.Tensor, metadata: Dict = None):
+        super().__init__()
+        self._is_inplace_pinned = set()
+        assert isinstance(torch_feature, torch.Tensor), (
+            f"torch_feature in TorchBasedFeature must be torch.Tensor, "
+            f"but got {type(torch_feature)}."
+        )
+        assert torch_feature.dim() > 1, (
+            f"dimension of torch_feature in TorchBasedFeature must be greater "
+            f"than 1, but got {torch_feature.dim()} dimension."
+        )
+        # Make sure the tensor is contiguous.
+        self._tensor = torch_feature.contiguous()
+        self._metadata = metadata
+
+    def __del__(self):
+        # torch.Tensor.pin_memory() is not an inplace operation. To make it
+        # truly in-place, we need to use cudaHostRegister. Then, we need to use
+        # cudaHostUnregister to unpin the tensor in the destructor.
+        # https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
+        for tensor in self._is_inplace_pinned:
+            assert self._inplace_unpinner(tensor.data_ptr()) == 0
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+
+        If the feature is on pinned CPU memory and `ids` is on GPU or pinned CPU
+        memory, it will be read by GPU and the returned tensor will be on GPU.
+        Otherwise, the returned tensor will be on CPU.
+
+        Parameters
+        ----------
+        ids : torch.Tensor, optional
+            The index of the feature. If specified, only the specified indices
+            of the feature are read. If None, the entire feature is returned.
+
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            if self._tensor.is_pinned():
+                return self._tensor.cuda()
+            return self._tensor
+        return index_select(self._tensor, ids)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        assert self._tensor.device.type == "cpu"
+        if ids.is_cuda and self.is_pinned():
+            current_stream = torch.cuda.current_stream()
+            host_to_device_stream = get_host_to_device_uva_stream()
+            host_to_device_stream.wait_stream(current_stream)
+            with torch.cuda.stream(host_to_device_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                values = index_select(self._tensor, ids)
+                values.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values)
+        elif ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to(self._tensor.device, non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            values = torch.ops.graphbolt.index_select_async(self._tensor, ids)
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values_cuda = values.wait().to(ids_device, non_blocking=True)
+                values_cuda.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values_cuda)
+        else:
+            yield torch.ops.graphbolt.index_select_async(self._tensor, ids)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        if ids_device.type == "cuda":
+            if self._tensor.is_cuda:
+                # If the ids and the tensor are on cuda, no need for async.
+                return 0
+            return 1 if self.is_pinned() else 3
+        else:
+            return 1
+
+    def size(self):
+        """Get the size of the feature.
+
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._tensor.size()[1:]
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._tensor.size()[0]
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Update the feature store.
+
+        Parameters
+        ----------
+        value : torch.Tensor
+            The updated value of the feature.
+        ids : torch.Tensor, optional
+            The indices of the feature to update. If specified, only the
+            specified indices of the feature will be updated. For the feature,
+            the `ids[i]` row is updated to `value[i]`. So the indices and value
+            must have the same length. If None, the entire feature will be
+            updated.
+        """
+        if ids is None:
+            self._tensor = value
+        else:
+            assert ids.shape[0] == value.shape[0], (
+                f"ids and value must have the same length, "
+                f"but got {ids.shape[0]} and {value.shape[0]}."
+            )
+            assert self.size() == value.size()[1:], (
+                f"The size of the feature is {self.size()}, "
+                f"while the size of the value is {value.size()[1:]}."
+            )
+            if self._tensor.is_pinned() and value.is_cuda and ids.is_cuda:
+                raise NotImplementedError(
+                    "Update the feature on pinned CPU memory by GPU is not "
+                    "supported yet."
+                )
+            self._tensor[ids] = value
+
+    def metadata(self):
+        """Get the metadata of the feature.
+
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return (
+            self._metadata if self._metadata is not None else super().metadata()
+        )
+
+    def pin_memory_(self):
+        """In-place operation to copy the feature to pinned memory. Returns the
+        same object modified in-place."""
+        # torch.Tensor.pin_memory() is not an inplace operation. To make it
+        # truly in-place, we need to use cudaHostRegister. Then, we need to use
+        # cudaHostUnregister to unpin the tensor in the destructor.
+        # https://github.com/pytorch/pytorch/issues/32167#issuecomment-753551842
+        x = self._tensor
+        if not x.is_pinned() and x.device.type == "cpu":
+            assert (
+                x.is_contiguous()
+            ), "Tensor pinning is only supported for contiguous tensors."
+            cudart = torch.cuda.cudart()
+            assert (
+                cudart.cudaHostRegister(
+                    x.data_ptr(), x.numel() * x.element_size(), 0
+                )
+                == 0
+            )
+
+            self._is_inplace_pinned.add(x)
+            self._inplace_unpinner = cudart.cudaHostUnregister
+
+        return self
+
+    def is_pinned(self):
+        """Returns True if the stored feature is pinned."""
+        return self._tensor.is_pinned()
+
+    def to(self, device):  # pylint: disable=invalid-name
+        """Copy `TorchBasedFeature` to the specified device."""
+        # copy.copy is a shallow copy so it does not copy tensor memory.
+        self2 = copy.copy(self)
+        if device == "pinned":
+            self2._tensor = self2._tensor.pin_memory()
+        else:
+            self2._tensor = self2._tensor.to(device)
+        return self2
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    feature={feature},\n"
+            "    metadata={metadata},\n"
+            ")"
+        )
+
+        feature_str = textwrap.indent(
+            str(self._tensor), " " * len("    feature=")
+        ).strip()
+        metadata_str = textwrap.indent(
+            str(self.metadata()), " " * len("    metadata=")
+        ).strip()
+
+        return ret.format(
+            Classname=self.__class__.__name__,
+            feature=feature_str,
+            metadata=metadata_str,
+        )
+
+
+class DiskBasedFeature(Feature):
+    r"""A wrapper of disk based feature.
+
+    Initialize a disk based feature fetcher by a numpy file. Note that you can
+    use gb.numpy_save_aligned as a replacement for np.save to potentially get
+    increased performance.
+
+    Parameters
+    ----------
+    path : string
+        The path to the numpy feature file.
+        Note that the dimension of the numpy should be greater than 1.
+    metadata : Dict
+        The metadata of the feature.
+    num_threads : int
+        The number of threads driving io_uring queues.
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> torch_feat = torch.arange(10).reshape(2, -1)
+    >>> pth = "path/to/feat.npy"
+    >>> np.save(pth, torch_feat)
+    >>> feature = gb.DiskBasedFeature(pth)
+    >>> feature.read(torch.tensor([0]))
+    tensor([[0, 1, 2, 3, 4]])
+    >>> feature.size()
+    torch.Size([5])
+    """
+
+    def __init__(self, path: str, metadata: Dict = None, num_threads=None):
+        super().__init__()
+        mmap_mode = "r+"
+        ondisk_data = np.load(path, mmap_mode=mmap_mode)
+        assert ondisk_data.flags[
+            "C_CONTIGUOUS"
+        ], "DiskBasedFeature only supports C_CONTIGUOUS array."
+        self._tensor = torch.from_numpy(ondisk_data)
+
+        self._metadata = metadata
+        if torch.ops.graphbolt.detect_io_uring():
+            self._ondisk_npy_array = torch.ops.graphbolt.ondisk_npy_array(
+                path, self._tensor.dtype, self._tensor.shape, num_threads
+            )
+
+    def read(self, ids: torch.Tensor = None):
+        """Read the feature by index.
+        The returned tensor will be on CPU.
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        torch.Tensor
+            The read feature.
+        """
+        if ids is None:
+            return self._tensor
+        elif torch.ops.graphbolt.detect_io_uring():
+            try:
+                return self._ondisk_npy_array.index_select(ids).wait()
+            except RuntimeError:
+                raise IndexError
+        else:
+            return index_select(self._tensor, ids)
+
+    def read_async(self, ids: torch.Tensor):
+        r"""Read the feature by index asynchronously.
+
+        Parameters
+        ----------
+        ids : torch.Tensor
+            The index of the feature. Only the specified indices of the
+            feature are read.
+        Returns
+        -------
+        A generator object.
+            The returned generator object returns a future on
+            ``read_async_num_stages(ids.device)``\ th invocation. The return result
+            can be accessed by calling ``.wait()``. on the returned future object.
+            It is undefined behavior to call ``.wait()`` more than once.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> feature = gb.Feature(...)
+        >>> ids = torch.tensor([0, 2])
+        >>> for stage, future in enumerate(feature.read_async(ids)):
+        ...     pass
+        >>> assert stage + 1 == feature.read_async_num_stages(ids.device)
+        >>> result = future.wait()  # result contains the read values.
+        """
+        assert torch.ops.graphbolt.detect_io_uring()
+        if ids.is_cuda:
+            ids_device = ids.device
+            current_stream = torch.cuda.current_stream()
+            device_to_host_stream = get_device_to_host_uva_stream()
+            device_to_host_stream.wait_stream(current_stream)
+            with torch.cuda.stream(device_to_host_stream):
+                ids.record_stream(torch.cuda.current_stream())
+                ids = ids.to(self._tensor.device, non_blocking=True)
+                ids_copy_event = torch.cuda.Event()
+                ids_copy_event.record()
+
+            yield  # first stage is done.
+
+            ids_copy_event.synchronize()
+            values = self._ondisk_npy_array.index_select(ids)
+            yield
+
+            host_to_device_stream = get_host_to_device_uva_stream()
+            with torch.cuda.stream(host_to_device_stream):
+                values_cuda = values.wait().to(ids_device, non_blocking=True)
+                values_cuda.record_stream(current_stream)
+                values_copy_event = torch.cuda.Event()
+                values_copy_event.record()
+
+            yield _Waiter(values_copy_event, values_cuda)
+        else:
+            yield self._ondisk_npy_array.index_select(ids)
+
+    def read_async_num_stages(self, ids_device: torch.device):
+        """The number of stages of the read_async operation. See read_async
+        function for directions on its use. This function is required to return
+        the number of yield operations when read_async is used with a tensor
+        residing on ids_device.
+
+        Parameters
+        ----------
+        ids_device : torch.device
+            The device of the ids parameter passed into read_async.
+        Returns
+        -------
+        int
+            The number of stages of the read_async operation.
+        """
+        return 3 if ids_device.type == "cuda" else 1
+
+    def size(self):
+        """Get the size of the feature.
+        Returns
+        -------
+        torch.Size
+            The size of the feature.
+        """
+        return self._tensor.size()[1:]
+
+    def count(self):
+        """Get the count of the feature.
+
+        Returns
+        -------
+        int
+            The count of the feature.
+        """
+        return self._tensor.size()[0]
+
+    def update(self, value: torch.Tensor, ids: torch.Tensor = None):
+        """Disk based feature does not support update for now."""
+        raise NotImplementedError
+
+    def metadata(self):
+        """Get the metadata of the feature.
+        Returns
+        -------
+        Dict
+            The metadata of the feature.
+        """
+        return (
+            self._metadata if self._metadata is not None else super().metadata()
+        )
+
+    def read_into_memory(self) -> TorchBasedFeature:
+        """Change disk-based feature to torch-based feature."""
+        return TorchBasedFeature(self._tensor, self._metadata)
+
+    def to(self, _):  # pylint: disable=invalid-name
+        """Placeholder `DiskBasedFeature` to implementation. It is a no-op."""
+        gb_warning(
+            "`DiskBasedFeature.to(device)` is not supported. Leaving unmodified."
+        )
+        return self
+
+    def pin_memory_(self):  # pylint: disable=invalid-name
+        r"""Placeholder `DiskBasedFeature` pin_memory_ implementation. It is a no-op."""
+        gb_warning(
+            "`DiskBasedFeature.pin_memory_()` is not supported. Leaving unmodified."
+        )
+        return self
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    feature={feature},\n"
+            "    metadata={metadata},\n"
+            ")"
+        )
+
+        feature_str = textwrap.indent(
+            str(self._tensor), " " * len("    feature=")
+        ).strip()
+        metadata_str = textwrap.indent(
+            str(self.metadata()), " " * len("    metadata=")
+        ).strip()
+
+        return ret.format(
+            Classname=self.__class__.__name__,
+            feature=feature_str,
+            metadata=metadata_str,
+        )
+
+
+class TorchBasedFeatureStore(BasicFeatureStore):
+    r"""A store to manage multiple pytorch based feature for access.
+
+    The feature stores are described by the `feat_data`. The `feat_data` is a
+    list of `OnDiskFeatureData`.
+
+    For a feature store, its format must be either "pt" or "npy" for Pytorch or
+    Numpy formats. If the format is "pt", the feature store must be loaded in
+    memory. If the format is "npy", the feature store can be loaded in memory or
+    on disk. Note that you can use gb.numpy_save_aligned as a replacement for
+    np.save to potentially get increased performance.
+
+    Parameters
+    ----------
+    feat_data : List[OnDiskFeatureData]
+        The description of the feature stores.
+
+    Examples
+    --------
+    >>> import torch
+    >>> import numpy as np
+    >>> from dgl import graphbolt as gb
+    >>> edge_label = torch.tensor([[1], [2], [3]])
+    >>> node_feat = torch.tensor([[1, 2, 3], [4, 5, 6]])
+    >>> torch.save(edge_label, "/tmp/edge_label.pt")
+    >>> gb.numpy_save_aligned("/tmp/node_feat.npy", node_feat.numpy())
+    >>> feat_data = [
+    ...     gb.OnDiskFeatureData(domain="edge", type="author:writes:paper",
+    ...         name="label", format="torch", path="/tmp/edge_label.pt",
+    ...         in_memory=True),
+    ...     gb.OnDiskFeatureData(domain="node", type="paper", name="feat",
+    ...         format="numpy", path="/tmp/node_feat.npy", in_memory=False),
+    ... ]
+    >>> feature_store = gb.TorchBasedFeatureStore(feat_data)
+    """
+
+    def __init__(self, feat_data: List[OnDiskFeatureData]):
+        features = {}
+        for spec in feat_data:
+            key = (spec.domain, spec.type, spec.name)
+            metadata = spec.extra_fields
+            if spec.format == "torch":
+                assert spec.in_memory, (
+                    f"Pytorch tensor can only be loaded in memory, "
+                    f"but the feature {key} is loaded on disk."
+                )
+                features[key] = TorchBasedFeature(
+                    torch.load(spec.path, weights_only=False), metadata=metadata
+                )
+            elif spec.format == "numpy":
+                if spec.in_memory:
+                    # TorchBasedFeature is always in memory by default.
+                    features[key] = TorchBasedFeature(
+                        torch.as_tensor(np.load(spec.path)), metadata=metadata
+                    )
+                else:
+                    # DiskBasedFeature is always out of memory by default.
+                    features[key] = DiskBasedFeature(
+                        spec.path, metadata=metadata
+                    )
+            else:
+                raise ValueError(f"Unknown feature format {spec.format}")
+        super().__init__(features)
+
+    def pin_memory_(self):
+        """In-place operation to copy the feature store to pinned memory.
+        Returns the same object modified in-place."""
+        if is_wsl():
+            gb_warning(
+                "In place pinning is not supported on WSL. "
+                "Returning the out of place pinned `TorchBasedFeatureStore`."
+            )
+            return self.to("pinned")
+        for feature in self._features.values():
+            feature.pin_memory_()
+
+        return self
+
+    def is_pinned(self):
+        """Returns True if all the stored features are pinned."""
+        return all(feature.is_pinned() for feature in self._features.values())
+
+    def to(self, device):  # pylint: disable=invalid-name
+        """Copy `TorchBasedFeatureStore` to the specified device."""
+        # copy.copy is a shallow copy so it does not copy tensor memory.
+        self2 = copy.copy(self)
+        self2._features = {k: v.to(device) for k, v in self2._features.items()}
+        return self2
+
+    def __repr__(self) -> str:
+        ret = "{Classname}(\n" + "    {features}\n" + ")"
+        features_str = textwrap.indent(str(self._features), "    ").strip()
+        return ret.format(
+            Classname=self.__class__.__name__, features=features_str
+        )

python/dgl/graphbolt/impl/uniform_negative_sampler.py

@@ -0,0 +1,86 @@
+"""Uniform negative sampler for GraphBolt."""
+
+import torch
+from torch.utils.data import functional_datapipe
+
+from ..negative_sampler import NegativeSampler
+
+__all__ = ["UniformNegativeSampler"]
+
+
+@functional_datapipe("sample_uniform_negative")
+class UniformNegativeSampler(NegativeSampler):
+    """Sample negative destination nodes for each source node based on a uniform
+    distribution.
+
+    Functional name: :obj:`sample_uniform_negative`.
+
+    It's important to note that the term 'negative' refers to false negatives,
+    indicating that the sampled pairs are not ensured to be absent in the graph.
+    For each edge ``(u, v)``, it is supposed to generate `negative_ratio` pairs
+    of negative edges ``(u, v')``, where ``v'`` is chosen uniformly from all
+    the nodes in the graph.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    graph : FusedCSCSamplingGraph
+        The graph on which to perform negative sampling.
+    negative_ratio : int
+        The proportion of negative samples to positive samples.
+
+    Examples
+    --------
+    >>> from dgl import graphbolt as gb
+    >>> indptr = torch.LongTensor([0, 1, 2, 3, 4])
+    >>> indices = torch.LongTensor([1, 2, 3, 0])
+    >>> graph = gb.fused_csc_sampling_graph(indptr, indices)
+    >>> seeds = torch.tensor([[0, 1], [1, 2], [2, 3], [3, 0]])
+    >>> item_set = gb.ItemSet(seeds, names="seeds")
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4,)
+    >>> neg_sampler = gb.UniformNegativeSampler(
+    ...     item_sampler, graph, 2)
+    >>> for minibatch in neg_sampler:
+    ...       print(minibatch.seeds)
+    ...       print(minibatch.labels)
+    ...       print(minibatch.indexes)
+    tensor([[0, 1], [1, 2], [2, 3], [3, 0], [0, 1], [0, 3], [1, 1], [1, 2],
+        [2, 1], [2, 0], [3, 0], [3, 2]])
+    tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.])
+    tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3])
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        graph,
+        negative_ratio,
+    ):
+        super().__init__(datapipe, negative_ratio)
+        self.graph = graph
+
+    def _sample_with_etype(self, seeds, etype=None):
+        assert seeds.ndim == 2 and seeds.shape[1] == 2, (
+            "Only tensor with shape N*2 is supported for negative"
+            + f" sampling, but got {seeds.shape}."
+        )
+        # Sample negative edges, and concatenate positive edges with them.
+        all_seeds = self.graph.sample_negative_edges_uniform(
+            etype,
+            seeds,
+            self.negative_ratio,
+        )
+        # Construct indexes for all node pairs.
+        pos_num = seeds.shape[0]
+        negative_ratio = self.negative_ratio
+        pos_indexes = torch.arange(0, pos_num, device=all_seeds.device)
+        neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
+        indexes = torch.cat((pos_indexes, neg_indexes))
+        # Construct labels for all node pairs.
+        neg_num = all_seeds.shape[0] - pos_num
+        labels = torch.empty(pos_num + neg_num, device=all_seeds.device)
+        labels[:pos_num] = 1
+        labels[pos_num:] = 0
+        return all_seeds, labels, indexes

python/dgl/graphbolt/internal/__init__.py

@@ -0,0 +1,4 @@
+"""Utility functions for GraphBolt."""
+from .utils import *
+from .sample_utils import *
+from .item_sampler_utils import *

python/dgl/graphbolt/internal/item_sampler_utils.py

@@ -0,0 +1,112 @@
+"""Utility functions for DistributedItemSampler."""
+
+
+def count_split(total, num_workers, worker_id, batch_size=1):
+    """Calculate the number of assigned items after splitting them by batch
+    size evenly. It will return the number for this worker and also a sum of
+    previous workers.
+    """
+    quotient, remainder = divmod(total, num_workers * batch_size)
+    if batch_size == 1:
+        assigned = quotient + (worker_id < remainder)
+    else:
+        batch_count, last_batch = divmod(remainder, batch_size)
+        assigned = quotient * batch_size + (
+            batch_size
+            if worker_id < batch_count
+            else (last_batch if worker_id == batch_count else 0)
+        )
+    prefix_sum = quotient * worker_id * batch_size + min(
+        worker_id * batch_size, remainder
+    )
+    return (assigned, prefix_sum)
+
+
+def calculate_range(
+    distributed,
+    total,
+    num_replicas,
+    rank,
+    num_workers,
+    worker_id,
+    batch_size,
+    drop_last,
+    drop_uneven_inputs,
+):
+    """Calculates the range of items to be assigned to the current worker.
+
+    This function evenly distributes `total` items among multiple workers,
+    batching them using `batch_size`. Each replica has `num_workers` workers.
+    The batches generated by workers within the same replica are combined into
+    the replica`s output. The `drop_last` parameter determines whether
+    incomplete batches should be dropped. If `drop_last` is True, incomplete
+    batches are discarded. The `drop_uneven_inputs` parameter determines if the
+    number of batches assigned to each replica should be the same. If
+    `drop_uneven_inputs` is True, excessive batches for some replicas will be
+    dropped.
+
+    Args:
+        distributed (bool): Whether it's in distributed mode.
+        total (int): The total number of items.
+        num_replicas (int): The total number of replicas.
+        rank (int): The rank of the current replica.
+        num_workers (int): The number of workers per replica.
+        worker_id (int): The ID of the current worker.
+        batch_size (int): The desired batch size.
+        drop_last (bool): Whether to drop incomplete batches.
+        drop_uneven_inputs (bool): Whether to drop excessive batches for some
+          replicas.
+
+    Returns:
+        tuple: A tuple containing three numbers:
+            - start_offset (int): The starting offset of the range assigned to
+              the current worker.
+            - assigned_count (int): The length of the range assigned to the
+              current worker.
+            - output_count (int): The number of items that the current worker
+              will produce after dropping.
+    """
+    # Check if it's distributed mode.
+    if not distributed:
+        if not drop_last:
+            return (0, total, total)
+        else:
+            return (0, total, total // batch_size * batch_size)
+    # First, equally distribute items into all replicas.
+    assigned_count, start_offset = count_split(
+        total, num_replicas, rank, batch_size
+    )
+    # Calculate the number of outputs when drop_uneven_inputs is True.
+    # `assigned_count` is the number of items distributed to the current
+    # process. `output_count` is the number of items should be output
+    # by this process after dropping.
+    if not drop_uneven_inputs:
+        if not drop_last:
+            output_count = assigned_count
+        else:
+            output_count = assigned_count // batch_size * batch_size
+    else:
+        if not drop_last:
+            min_item_count, _ = count_split(
+                total, num_replicas, num_replicas - 1, batch_size
+            )
+            min_batch_count = (min_item_count + batch_size - 1) // batch_size
+            output_count = min(min_batch_count * batch_size, assigned_count)
+        else:
+            output_count = total // (batch_size * num_replicas) * batch_size
+    # If there are multiple workers, equally distribute the batches to
+    # all workers.
+    if num_workers > 1:
+        # Equally distribute the dropped number too.
+        dropped_items, prev_dropped_items = count_split(
+            assigned_count - output_count, num_workers, worker_id
+        )
+        output_count, prev_output_count = count_split(
+            output_count,
+            num_workers,
+            worker_id,
+            batch_size,
+        )
+        assigned_count = output_count + dropped_items
+        start_offset += prev_output_count + prev_dropped_items
+    return (start_offset, assigned_count, output_count)

python/dgl/graphbolt/internal/sample_utils.py

@@ -0,0 +1,568 @@
+"""Utility functions for sampling."""
+
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+from ..base import CSCFormatBase, etype_str_to_tuple, expand_indptr
+
+
+def unique_and_compact(
+    nodes: Union[
+        List[torch.Tensor],
+        Dict[str, List[torch.Tensor]],
+    ],
+    rank: int = 0,
+    world_size: int = 1,
+    async_op: bool = False,
+):
+    """
+    Compact a list of nodes tensor. The `rank` and `world_size` parameters are
+    relevant when using Cooperative Minibatching, which was initially proposed
+    in `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
+    was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching.
+
+    When `world_size` is greater than 1, then the given ids are partitioned
+    between the available ranks. The ids corresponding to the given rank are
+    guaranteed to come before the ids of other ranks. To do this, the
+    partitioned ids are rotated backwards by the given rank so that the ids are
+    ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
+    supported only for Volta and later generation NVIDIA GPUs.
+
+    Parameters
+    ----------
+    nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
+        List of nodes for compacting.
+        the unique_and_compact will be done per type
+        - If `nodes` is a list of tensor: All the tensors will do unique and
+        compact together, usually it is used for homogeneous graph.
+        - If `nodes` is a list of dictionary: The keys should be node type and
+        the values should be corresponding nodes, the unique and compact will
+        be done per type, usually it is used for heterogeneous graph.
+    rank : int
+        The rank of the current process.
+    world_size : int
+        The number of processes.
+    async_op: bool
+        Boolean indicating whether the call is asynchronous. If so, the result
+        can be obtained by calling wait on the returned future.
+
+    Returns
+    -------
+    Tuple[unique_nodes, compacted_node_list, unique_nodes_offsets]
+        The Unique nodes (per type) of all nodes in the input. And the compacted
+        nodes list, where IDs inside are replaced with compacted node IDs.
+        "Compacted node list" indicates that the node IDs in the input node
+        list are replaced with mapped node IDs, where each type of node is
+        mapped to a contiguous space of IDs ranging from 0 to N.
+        The unique nodes offsets tensor partitions the unique_nodes tensor. Has
+        size `world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]`
+        belongs to the rank `(rank + i) % world_size`.
+    """
+    is_heterogeneous = isinstance(nodes, dict)
+
+    if not is_heterogeneous:
+        homo_ntype = "a"
+        nodes = {homo_ntype: nodes}
+
+    nums = {}
+    concat_nodes, empties = [], []
+    for ntype, nodes_of_type in nodes.items():
+        nums[ntype] = [node.size(0) for node in nodes_of_type]
+        concat_nodes.append(torch.cat(nodes_of_type))
+        empties.append(concat_nodes[-1].new_empty(0))
+    unique_fn = (
+        torch.ops.graphbolt.unique_and_compact_batched_async
+        if async_op
+        else torch.ops.graphbolt.unique_and_compact_batched
+    )
+    results = unique_fn(concat_nodes, empties, empties, rank, world_size)
+
+    class _Waiter:
+        def __init__(self, future, ntypes, nums):
+            self.future = future
+            self.ntypes = ntypes
+            self.nums = nums
+
+        def wait(self):
+            """Returns the stored value when invoked."""
+            results = self.future.wait() if async_op else self.future
+            ntypes = self.ntypes
+            nums = self.nums
+            # Ensure there is no memory leak.
+            self.future = self.ntypes = self.nums = None
+
+            unique, compacted, offsets = {}, {}, {}
+            for ntype, result in zip(ntypes, results):
+                (
+                    unique[ntype],
+                    concat_compacted,
+                    _,
+                    offsets[ntype],
+                ) = result
+                compacted[ntype] = list(concat_compacted.split(nums[ntype]))
+            if is_heterogeneous:
+                return unique, compacted, offsets
+            else:
+                return (
+                    unique[homo_ntype],
+                    compacted[homo_ntype],
+                    offsets[homo_ntype],
+                )
+
+    post_processer = _Waiter(results, nodes.keys(), nums)
+    if async_op:
+        return post_processer
+    else:
+        return post_processer.wait()
+
+
+def compact_temporal_nodes(nodes, nodes_timestamp):
+    """Compact a list of temporal nodes without unique.
+
+    Note that since there is no unique, the nodes and nodes_timestamp are simply
+    concatenated. And the compacted nodes are consecutive numbers starting from
+    0.
+
+    Parameters
+    ----------
+    nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
+        List of nodes for compacting.
+        the compact operator will be done per type
+        - If `nodes` is a list of tensor: All the tensors will compact together,
+        usually it is used for homogeneous graph.
+        - If `nodes` is a list of dictionary: The keys should be node type and
+        the values should be corresponding nodes, the compact will be done per
+        type, usually it is used for heterogeneous graph.
+
+    nodes_timestamp : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
+        List of timestamps for compacting.
+
+    Returns
+    -------
+    Tuple[nodes, nodes_timestamp, compacted_node_list]
+
+    The concatenated nodes and nodes_timestamp, and the compacted nodes list,
+    where IDs inside are replaced with compacted node IDs.
+    """
+
+    def _compact_per_type(per_type_nodes, per_type_nodes_timestamp):
+        nums = [node.size(0) for node in per_type_nodes]
+        per_type_nodes = torch.cat(per_type_nodes)
+        per_type_nodes_timestamp = torch.cat(per_type_nodes_timestamp)
+        compacted_nodes = torch.arange(
+            0,
+            per_type_nodes.numel(),
+            dtype=per_type_nodes.dtype,
+            device=per_type_nodes.device,
+        )
+        compacted_nodes = list(compacted_nodes.split(nums))
+        return per_type_nodes, per_type_nodes_timestamp, compacted_nodes
+
+    if isinstance(nodes, dict):
+        ret_nodes, ret_timestamp, compacted = {}, {}, {}
+        for ntype, nodes_of_type in nodes.items():
+            (
+                ret_nodes[ntype],
+                ret_timestamp[ntype],
+                compacted[ntype],
+            ) = _compact_per_type(nodes_of_type, nodes_timestamp[ntype])
+        return ret_nodes, ret_timestamp, compacted
+    else:
+        return _compact_per_type(nodes, nodes_timestamp)
+
+
+def unique_and_compact_csc_formats(
+    csc_formats: Union[
+        Tuple[torch.Tensor, torch.Tensor],
+        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
+    ],
+    unique_dst_nodes: Union[
+        torch.Tensor,
+        Dict[str, torch.Tensor],
+    ],
+    rank: int = 0,
+    world_size: int = 1,
+    async_op: bool = False,
+):
+    """
+    Compact csc formats and return unique nodes (per type). The `rank` and
+    `world_size` parameters are relevant when using Cooperative Minibatching,
+    which was initially proposed in
+    `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__
+    and was later first fully described in
+    `Cooperative Minibatching in Graph Neural Networks
+    <https://arxiv.org/abs/2310.12403>`__.
+    Cooperation between the GPUs eliminates duplicate work performed across the
+    GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
+    performing GNN minibatching.
+
+    When `world_size` is greater than 1, then the given ids are partitioned
+    between the available ranks. The ids corresponding to the given rank are
+    guaranteed to come before the ids of other ranks. To do this, the
+    partitioned ids are rotated backwards by the given rank so that the ids are
+    ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
+    supported only for Volta and later generation NVIDIA GPUs.
+
+    Parameters
+    ----------
+    csc_formats : Union[CSCFormatBase, Dict(str, CSCFormatBase)]
+        CSC formats representing source-destination edges.
+        - If `csc_formats` is a CSCFormatBase: It means the graph is
+        homogeneous. Also, indptr and indice in it should be torch.tensor
+        representing source and destination pairs in csc format. And IDs inside
+        are homogeneous ids.
+        - If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
+        should be edge type and the values should be csc format node pairs.
+        And IDs inside are heterogeneous ids.
+    unique_dst_nodes: torch.Tensor or Dict[str, torch.Tensor]
+        Unique nodes of all destination nodes in the node pairs.
+        - If `unique_dst_nodes` is a tensor: It means the graph is homogeneous.
+        - If `csc_formats` is a dictionary: The keys are node type and the
+        values are corresponding nodes. And IDs inside are heterogeneous ids.
+    rank : int
+        The rank of the current process.
+    world_size : int
+        The number of processes.
+    async_op: bool
+        Boolean indicating whether the call is asynchronous. If so, the result
+        can be obtained by calling wait on the returned future.
+
+    Returns
+    -------
+    Tuple[unique_nodes, csc_formats, unique_nodes_offsets]
+        The compacted csc formats, where node IDs are replaced with mapped node
+        IDs, and the unique nodes (per type).
+        "Compacted csc formats" indicates that the node IDs in the input node
+        pairs are replaced with mapped node IDs, where each type of node is
+        mapped to a contiguous space of IDs ranging from 0 to N. The unique
+        nodes offsets tensor partitions the unique_nodes tensor. Has size
+        `world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]` belongs
+        to the rank `(rank + i) % world_size`.
+
+    Examples
+    --------
+    >>> import dgl.graphbolt as gb
+    >>> N1 = torch.LongTensor([1, 2, 2])
+    >>> N2 = torch.LongTensor([5, 5, 6])
+    >>> unique_dst = {
+    ...     "n1": torch.LongTensor([1, 2]),
+    ...     "n2": torch.LongTensor([5, 6])}
+    >>> csc_formats = {
+    ...     "n1:e1:n2": gb.CSCFormatBase(indptr=torch.tensor([0, 2, 3]),indices=N1),
+    ...     "n2:e2:n1": gb.CSCFormatBase(indptr=torch.tensor([0, 1, 3]),indices=N2)}
+    >>> unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
+    ...     csc_formats, unique_dst
+    ... )
+    >>> print(unique_nodes)
+    {'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
+    >>> print(compacted_csc_formats)
+    {"n1:e1:n2": CSCFormatBase(indptr=torch.tensor([0, 2, 3]),
+                               indices=torch.tensor([0, 1, 1])),
+     "n2:e2:n1": CSCFormatBase(indptr=torch.tensor([0, 1, 3]),
+                               indices=torch.Longtensor([0, 0, 1]))}
+    """
+    is_homogeneous = not isinstance(csc_formats, dict)
+    if is_homogeneous:
+        csc_formats = {"_N:_E:_N": csc_formats}
+        if unique_dst_nodes is not None:
+            assert isinstance(
+                unique_dst_nodes, torch.Tensor
+            ), "Edge type not supported in homogeneous graph."
+            unique_dst_nodes = {"_N": unique_dst_nodes}
+
+    # Collect all source and destination nodes for each node type.
+    indices = defaultdict(list)
+    device = None
+    for etype, csc_format in csc_formats.items():
+        if device is None:
+            device = csc_format.indices.device
+        src_type, _, dst_type = etype_str_to_tuple(etype)
+        assert len(unique_dst_nodes.get(dst_type, [])) + 1 == len(
+            csc_format.indptr
+        ), "The seed nodes should correspond to indptr."
+        indices[src_type].append(csc_format.indices)
+    indices = {ntype: torch.cat(nodes) for ntype, nodes in indices.items()}
+
+    ntypes = set(indices.keys())
+    dtype = list(indices.values())[0].dtype
+    default_tensor = torch.tensor([], dtype=dtype, device=device)
+    indice_list = []
+    unique_dst_list = []
+    for ntype in ntypes:
+        indice_list.append(indices.get(ntype, default_tensor))
+        unique_dst_list.append(unique_dst_nodes.get(ntype, default_tensor))
+    dst_list = [torch.tensor([], dtype=dtype, device=device)] * len(
+        unique_dst_list
+    )
+    uniq_fn = (
+        torch.ops.graphbolt.unique_and_compact_batched_async
+        if async_op
+        else torch.ops.graphbolt.unique_and_compact_batched
+    )
+    results = uniq_fn(indice_list, dst_list, unique_dst_list, rank, world_size)
+
+    class _Waiter:
+        def __init__(self, future, csc_formats):
+            self.future = future
+            self.csc_formats = csc_formats
+
+        def wait(self):
+            """Returns the stored value when invoked."""
+            results = self.future.wait() if async_op else self.future
+            csc_formats = self.csc_formats
+            # Ensure there is no memory leak.
+            self.future = self.csc_formats = None
+
+            unique_nodes = {}
+            compacted_indices = {}
+            offsets = {}
+            for i, ntype in enumerate(ntypes):
+                (
+                    unique_nodes[ntype],
+                    compacted_indices[ntype],
+                    _,
+                    offsets[ntype],
+                ) = results[i]
+
+            compacted_csc_formats = {}
+            # Map back with the same order.
+            for etype, csc_format in csc_formats.items():
+                num_elem = csc_format.indices.size(0)
+                src_type, _, _ = etype_str_to_tuple(etype)
+                indice = compacted_indices[src_type][:num_elem]
+                indptr = csc_format.indptr
+                compacted_csc_formats[etype] = CSCFormatBase(
+                    indptr=indptr, indices=indice
+                )
+                compacted_indices[src_type] = compacted_indices[src_type][
+                    num_elem:
+                ]
+
+            # Return singleton for a homogeneous graph.
+            if is_homogeneous:
+                compacted_csc_formats = list(compacted_csc_formats.values())[0]
+                unique_nodes = list(unique_nodes.values())[0]
+                offsets = list(offsets.values())[0]
+
+            return unique_nodes, compacted_csc_formats, offsets
+
+    post_processer = _Waiter(results, csc_formats)
+    if async_op:
+        return post_processer
+    else:
+        return post_processer.wait()
+
+
+def _broadcast_timestamps(csc, dst_timestamps):
+    """Broadcast the timestamp of each destination node to its corresponding
+    source nodes."""
+    return expand_indptr(
+        csc.indptr, node_ids=dst_timestamps, output_size=len(csc.indices)
+    )
+
+
+def compact_csc_format(
+    csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]],
+    dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]],
+    dst_timestamps: Optional[
+        Union[torch.Tensor, Dict[str, torch.Tensor]]
+    ] = None,
+):
+    """
+    Relabel the row (source) IDs in the csc formats into a contiguous range from
+    0 and return the original row node IDs per type.
+
+    Note that
+    1. The column (destination) IDs are included in the relabeled row IDs.
+    2. If there are repeated row IDs, they would not be uniqued and will be
+    treated as different nodes.
+    3. If `dst_timestamps` is given, the timestamp of each destination node will
+    be broadcasted to its corresponding source nodes.
+
+    Parameters
+    ----------
+    csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]]
+        CSC formats representing source-destination edges.
+        - If `csc_formats` is a CSCFormatBase: It means the graph is
+        homogeneous. Also, indptr and indice in it should be torch.tensor
+        representing source and destination pairs in csc format. And IDs inside
+        are homogeneous ids.
+        - If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
+        should be edge type and the values should be csc format node pairs.
+        And IDs inside are heterogeneous ids.
+    dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]]
+        Nodes of all destination nodes in the node pairs.
+        - If `dst_nodes` is a tensor: It means the graph is homogeneous.
+        - If `dst_nodes` is a dictionary: The keys are node type and the
+        values are corresponding nodes. And IDs inside are heterogeneous ids.
+
+    dst_timestamps: Optional[Union[torch.Tensor, Dict[str, torch.Tensor]]]
+        Timestamps of all destination nodes in the csc formats.
+        If given, the timestamp of each destination node will be broadcasted
+        to its corresponding source nodes.
+
+    Returns
+    -------
+    Tuple[original_row_node_ids, compacted_csc_formats, ...]
+        A tensor of original row node IDs (per type) of all nodes in the input.
+        The compacted CSC formats, where node IDs are replaced with mapped node
+        IDs ranging from 0 to N.
+        The source timestamps (per type) of all nodes in the input if
+        `dst_timestamps` is given.
+
+    Examples
+    --------
+    >>> import dgl.graphbolt as gb
+    >>> csc_formats = {
+    ...     "n2:e2:n1": gb.CSCFormatBase(
+    ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
+    ...     ),
+    ...     "n1:e1:n1": gb.CSCFormatBase(
+    ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
+    ...     ),
+    ... }
+    >>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
+    >>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
+    ...     csc_formats, dst_nodes
+    ... )
+    >>> original_row_node_ids
+    {'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
+    >>> compacted_csc_formats
+    {'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
+                indices=tensor([0, 1, 2]),
+    ), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
+                indices=tensor([2, 3, 4]),
+    )}
+
+    >>> csc_formats = {
+    ...     "n2:e2:n1": gb.CSCFormatBase(
+    ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
+    ...     ),
+    ...     "n1:e1:n1": gb.CSCFormatBase(
+    ...         indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
+    ...     ),
+    ... }
+    >>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
+    >>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
+    ...     csc_formats, dst_nodes
+    ... )
+    >>> original_row_node_ids
+    {'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
+    >>> compacted_csc_formats
+    {'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
+                indices=tensor([0, 1, 2]),
+    ), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
+                indices=tensor([2, 3, 4]),
+    )}
+
+    >>> dst_timestamps = {"n1": torch.LongTensor([10, 20])}
+    >>> (
+    ...     original_row_node_ids,
+    ...     compacted_csc_formats,
+    ...     src_timestamps,
+    ... ) = gb.compact_csc_format(csc_formats, dst_nodes, dst_timestamps)
+    >>> src_timestamps
+    {'n1': tensor([10, 20, 10, 20, 20]), 'n2': tensor([10, 20, 20])}
+    """
+    is_homogeneous = not isinstance(csc_formats, dict)
+    has_timestamp = dst_timestamps is not None
+    if is_homogeneous:
+        if dst_nodes is not None:
+            assert isinstance(
+                dst_nodes, torch.Tensor
+            ), "Edge type not supported in homogeneous graph."
+            assert len(dst_nodes) + 1 == len(
+                csc_formats.indptr
+            ), "The seed nodes should correspond to indptr."
+        offset = dst_nodes.size(0)
+        original_row_ids = torch.cat((dst_nodes, csc_formats.indices))
+        compacted_csc_formats = CSCFormatBase(
+            indptr=csc_formats.indptr,
+            indices=(
+                torch.arange(
+                    0,
+                    csc_formats.indices.size(0),
+                    device=csc_formats.indices.device,
+                )
+                + offset
+            ),
+        )
+
+        src_timestamps = None
+        if has_timestamp:
+            src_timestamps = torch.cat(
+                [
+                    dst_timestamps,
+                    _broadcast_timestamps(
+                        compacted_csc_formats, dst_timestamps
+                    ),
+                ]
+            )
+    else:
+        compacted_csc_formats = {}
+        src_timestamps = None
+        original_row_ids = {key: val.clone() for key, val in dst_nodes.items()}
+        if has_timestamp:
+            src_timestamps = {
+                key: val.clone() for key, val in dst_timestamps.items()
+            }
+        for etype, csc_format in csc_formats.items():
+            src_type, _, dst_type = etype_str_to_tuple(etype)
+            assert len(dst_nodes.get(dst_type, [])) + 1 == len(
+                csc_format.indptr
+            ), "The seed nodes should correspond to indptr."
+            device = csc_format.indices.device
+            offset = original_row_ids.get(
+                src_type, torch.tensor([], device=device)
+            ).size(0)
+            original_row_ids[src_type] = torch.cat(
+                (
+                    original_row_ids.get(
+                        src_type,
+                        torch.tensor(
+                            [], dtype=csc_format.indices.dtype, device=device
+                        ),
+                    ),
+                    csc_format.indices,
+                )
+            )
+            compacted_csc_formats[etype] = CSCFormatBase(
+                indptr=csc_format.indptr,
+                indices=(
+                    torch.arange(
+                        0,
+                        csc_format.indices.size(0),
+                        dtype=csc_format.indices.dtype,
+                        device=device,
+                    )
+                    + offset
+                ),
+            )
+            if has_timestamp:
+                # If destination timestamps are given, broadcast them to the
+                # corresponding source nodes.
+                src_timestamps[src_type] = torch.cat(
+                    (
+                        src_timestamps.get(
+                            src_type,
+                            torch.tensor(
+                                [],
+                                dtype=dst_timestamps[dst_type].dtype,
+                                device=device,
+                            ),
+                        ),
+                        _broadcast_timestamps(
+                            csc_format, dst_timestamps[dst_type]
+                        ),
+                    )
+                )
+    if has_timestamp:
+        return original_row_ids, compacted_csc_formats, src_timestamps
+    return original_row_ids, compacted_csc_formats

python/dgl/graphbolt/internal/utils.py

@@ -0,0 +1,216 @@
+"""Utility functions for GraphBolt."""
+
+import hashlib
+import json
+import os
+import shutil
+from typing import List, Union
+
+import numpy as np
+import pandas as pd
+import torch
+from numpy.lib.format import read_array_header_1_0, read_array_header_2_0
+
+
+def numpy_save_aligned(*args, **kwargs):
+    """A wrapper for numpy.save(), ensures the array is stored 4KiB aligned."""
+    # https://github.com/numpy/numpy/blob/2093a6d5b933f812d15a3de0eafeeb23c61f948a/numpy/lib/format.py#L179
+    has_array_align = hasattr(np.lib.format, "ARRAY_ALIGN")
+    if has_array_align:
+        default_alignment = np.lib.format.ARRAY_ALIGN
+        # The maximum allowed alignment by the numpy code linked above is 4K.
+        # Most filesystems work with block sizes of 4K so in practice, the file
+        # size on the disk won't be larger.
+        np.lib.format.ARRAY_ALIGN = 4096
+    np.save(*args, **kwargs)
+    if has_array_align:
+        np.lib.format.ARRAY_ALIGN = default_alignment
+
+
+def _read_torch_data(path):
+    return torch.load(path, weights_only=False)
+
+
+def _read_numpy_data(path, in_memory=True):
+    if in_memory:
+        return torch.from_numpy(np.load(path))
+    return torch.as_tensor(np.load(path, mmap_mode="r+"))
+
+
+def read_data(path, fmt, in_memory=True):
+    """Read data from disk."""
+    if fmt == "torch":
+        return _read_torch_data(path)
+    elif fmt == "numpy":
+        return _read_numpy_data(path, in_memory=in_memory)
+    else:
+        raise RuntimeError(f"Unsupported format: {fmt}")
+
+
+def save_data(data, path, fmt):
+    """Save data into disk."""
+    # Make sure the directory exists.
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    if fmt not in ["numpy", "torch"]:
+        raise RuntimeError(f"Unsupported format: {fmt}")
+
+    # Perform necessary conversion.
+    if fmt == "numpy" and isinstance(data, torch.Tensor):
+        data = data.cpu().numpy()
+    elif fmt == "torch" and isinstance(data, np.ndarray):
+        data = torch.from_numpy(data).cpu()
+
+    # Save the data.
+    if fmt == "numpy":
+        if not data.flags["C_CONTIGUOUS"]:
+            Warning(
+                "The ndarray saved to disk is not contiguous, "
+                "so it will be copied to contiguous memory."
+            )
+            data = np.ascontiguousarray(data)
+        numpy_save_aligned(path, data)
+    elif fmt == "torch":
+        if not data.is_contiguous():
+            Warning(
+                "The tensor saved to disk is not contiguous, "
+                "so it will be copied to contiguous memory."
+            )
+            data = data.contiguous()
+        torch.save(data, path)
+
+
+def get_npy_dim(npy_path):
+    """Get the dim of numpy file."""
+    with open(npy_path, "rb") as f:
+        # For the read_array_header API provided by numpy will only read the
+        # length of the header, it will cause parsing failure and error if
+        # first 8 bytes which contains magin string and version are not read
+        # ahead of time. So, we need to make sure we have skipped these 8
+        # bytes.
+        f.seek(8, 0)
+        try:
+            shape, _, _ = read_array_header_1_0(f)
+        except ValueError:
+            try:
+                shape, _, _ = read_array_header_2_0(f)
+            except ValueError:
+                raise ValueError("Invalid file format")
+
+        return len(shape)
+
+
+def _to_int32(data):
+    if isinstance(data, torch.Tensor):
+        return data.to(torch.int32)
+    elif isinstance(data, np.ndarray):
+        return data.astype(np.int32)
+    else:
+        raise TypeError(
+            "Unsupported input type. Please provide a torch tensor or numpy array."
+        )
+
+
+def copy_or_convert_data(
+    input_path,
+    output_path,
+    input_format,
+    output_format="numpy",
+    in_memory=True,
+    is_feature=False,
+    within_int32=False,
+):
+    """Copy or convert the data from input_path to output_path."""
+    assert (
+        output_format == "numpy"
+    ), "The output format of the data should be numpy."
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    # We read the data always in case we need to cast its type.
+    data = read_data(input_path, input_format, in_memory)
+    if within_int32:
+        data = _to_int32(data)
+    if input_format == "numpy":
+        # If dim of the data is 1, reshape it to n * 1 and save it to output_path.
+        if is_feature and get_npy_dim(input_path) == 1:
+            data = data.reshape(-1, 1)
+        # If the data does not need to be modified, just copy the file.
+        elif not within_int32 and data.numpy().flags["C_CONTIGUOUS"]:
+            shutil.copyfile(input_path, output_path)
+            return
+    else:
+        # If dim of the data is 1, reshape it to n * 1 and save it to output_path.
+        if is_feature and data.dim() == 1:
+            data = data.reshape(-1, 1)
+    save_data(data, output_path, output_format)
+
+
+def read_edges(dataset_dir, edge_fmt, edge_path):
+    """Read egde data from numpy or csv."""
+    assert edge_fmt in [
+        "numpy",
+        "csv",
+    ], f"`numpy` or `csv` is expected when reading edges but got `{edge_fmt}`."
+    if edge_fmt == "numpy":
+        edge_data = read_data(
+            os.path.join(dataset_dir, edge_path),
+            edge_fmt,
+        )
+        assert (
+            edge_data.shape[0] == 2 and len(edge_data.shape) == 2
+        ), f"The shape of edges should be (2, N), but got {edge_data.shape}."
+        src, dst = edge_data.numpy()
+    else:
+        edge_data = pd.read_csv(
+            os.path.join(dataset_dir, edge_path),
+            names=["src", "dst"],
+        )
+        src, dst = edge_data["src"].to_numpy(), edge_data["dst"].to_numpy()
+    return (src, dst)
+
+
+def calculate_file_hash(file_path, hash_algo="md5"):
+    """Calculate the hash value of a file."""
+    hash_algos = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512"]
+    if hash_algo in hash_algos:
+        hash_obj = getattr(hashlib, hash_algo)()
+    else:
+        raise ValueError(
+            f"Hash algorithm must be one of: {hash_algos}, but got `{hash_algo}`."
+        )
+    with open(file_path, "rb") as file:
+        for chunk in iter(lambda: file.read(4096), b""):
+            hash_obj.update(chunk)
+    return hash_obj.hexdigest()
+
+
+def calculate_dir_hash(
+    dir_path, hash_algo="md5", ignore: Union[str, List[str]] = None
+):
+    """Calculte the hash values of all files under the directory."""
+    hashes = {}
+    for dirpath, _, filenames in os.walk(dir_path):
+        for filename in filenames:
+            if ignore and filename in ignore:
+                continue
+            filepath = os.path.join(dirpath, filename)
+            file_hash = calculate_file_hash(filepath, hash_algo=hash_algo)
+            hashes[filepath] = file_hash
+    return hashes
+
+
+def check_dataset_change(dataset_dir, processed_dir):
+    """Check whether dataset has been changed by checking its hash value."""
+    hash_value_file = "dataset_hash_value.txt"
+    hash_value_file_path = os.path.join(
+        dataset_dir, processed_dir, hash_value_file
+    )
+    if not os.path.exists(hash_value_file_path):
+        return True
+    with open(hash_value_file_path, "r") as f:
+        oringinal_hash_value = json.load(f)
+    present_hash_value = calculate_dir_hash(dataset_dir, ignore=hash_value_file)
+    if oringinal_hash_value == present_hash_value:
+        force_preprocess = False
+    else:
+        force_preprocess = True
+    return force_preprocess

python/dgl/graphbolt/internal_utils.py

@@ -0,0 +1,371 @@
+"""Miscallenous internal utils."""
+import functools
+import hashlib
+import os
+import platform
+import warnings
+from collections.abc import Mapping, Sequence
+
+import requests
+import torch
+from tqdm.auto import tqdm
+
+try:
+    from packaging import version  # pylint: disable=unused-import
+except ImportError:
+    # If packaging isn't installed, try and use the vendored copy in setuptools
+    from setuptools.extern.packaging import version
+
+
+@functools.lru_cache(maxsize=None)
+def is_wsl(v: str = platform.uname().release) -> int:
+    """Detects if Python is running in WSL"""
+
+    if v.endswith("-Microsoft"):
+        return 1
+    elif v.endswith("microsoft-standard-WSL2"):
+        return 2
+
+    return 0
+
+
+# pylint: disable=invalid-name
+_default_formatwarning = warnings.formatwarning
+
+
+def built_with_cuda():
+    """Returns whether GraphBolt was built with CUDA support."""
+    # This op is defined if graphbolt is built with CUDA support.
+    return hasattr(torch.ops.graphbolt, "set_max_uva_threads")
+
+
+class GBWarning(UserWarning):
+    """GraphBolt Warning class."""
+
+
+# pylint: disable=unused-argument
+def gb_warning_format(message, category, filename, lineno, line=None):
+    """Format GraphBolt warnings."""
+    if isinstance(category, GBWarning):
+        return "GraphBolt Warning: {}\n".format(message)
+    else:
+        return _default_formatwarning(
+            message, category, filename, lineno, line=None
+        )
+
+
+def gb_warning(message, category=GBWarning, stacklevel=2):
+    """GraphBolt warning wrapper that defaults to ``GBWarning`` instead of
+    ``UserWarning`` category.
+    """
+    return warnings.warn(message, category=category, stacklevel=stacklevel)
+
+
+warnings.formatwarning = gb_warning_format
+
+
+def is_listlike(data):
+    """Return if the data is a sequence but not a string."""
+    return isinstance(data, Sequence) and not isinstance(data, str)
+
+
+def recursive_apply(data, fn, *args, **kwargs):
+    """Recursively apply a function to every element in a container.
+
+    If the input data is a list or any sequence other than a string, returns a list
+    whose elements are the same elements applied with the given function.
+
+    If the input data is a dict or any mapping, returns a dict whose keys are the same
+    and values are the elements applied with the given function.
+
+    If the input data is a nested container, the result will have the same nested
+    structure where each element is transformed recursively.
+
+    The first argument of the function will be passed with the individual elements from
+    the input data, followed by the arguments in :attr:`args` and :attr:`kwargs`.
+
+    Parameters
+    ----------
+    data : any
+        Any object.
+    fn : callable
+        Any function.
+    args, kwargs :
+        Additional arguments and keyword-arguments passed to the function.
+
+    Examples
+    --------
+    Applying a ReLU function to a dictionary of tensors:
+
+    >>> h = {k: torch.randn(3) for k in ['A', 'B', 'C']}
+    >>> h = recursive_apply(h, torch.nn.functional.relu)
+    >>> assert all((v >= 0).all() for v in h.values())
+    """
+    if isinstance(data, Mapping):
+        return {
+            k: recursive_apply(v, fn, *args, **kwargs) for k, v in data.items()
+        }
+    elif isinstance(data, tuple):
+        return tuple(recursive_apply(v, fn, *args, **kwargs) for v in data)
+    elif is_listlike(data):
+        return [recursive_apply(v, fn, *args, **kwargs) for v in data]
+    else:
+        return fn(data, *args, **kwargs)
+
+
+def recursive_apply_reduce_all(data, fn, *args, **kwargs):
+    """Recursively apply a function to every element in a container and reduce
+    the boolean results with all.
+
+    If the input data is a list or any sequence other than a string, returns
+    True if and only if the given function returns True for all elements.
+
+    If the input data is a dict or any mapping, returns True if and only if the
+    given function returns True for values.
+
+    If the input data is a nested container, the result will be reduced over the
+    nested structure where each element is tested recursively.
+
+    The first argument of the function will be passed with the individual elements from
+    the input data, followed by the arguments in :attr:`args` and :attr:`kwargs`.
+
+    Parameters
+    ----------
+    data : any
+        Any object.
+    fn : callable
+        Any function returning a boolean.
+    args, kwargs :
+        Additional arguments and keyword-arguments passed to the function.
+    """
+    if isinstance(data, Mapping):
+        return all(
+            recursive_apply_reduce_all(v, fn, *args, **kwargs)
+            for v in data.values()
+        )
+    elif isinstance(data, tuple) or is_listlike(data):
+        return all(
+            recursive_apply_reduce_all(v, fn, *args, **kwargs) for v in data
+        )
+    else:
+        return fn(data, *args, **kwargs)
+
+
+def get_nonproperty_attributes(_obj) -> list:
+    """Get attributes of the class except for the properties."""
+    attributes = [
+        attribute
+        for attribute in dir(_obj)
+        if not attribute.startswith("__")
+        and (
+            not hasattr(type(_obj), attribute)
+            or not isinstance(getattr(type(_obj), attribute), property)
+        )
+        and not callable(getattr(_obj, attribute))
+    ]
+    return attributes
+
+
+def get_attributes(_obj) -> list:
+    """Get attributes of the class."""
+    attributes = [
+        attribute
+        for attribute in dir(_obj)
+        if not attribute.startswith("__")
+        and not callable(getattr(_obj, attribute))
+    ]
+    return attributes
+
+
+def download(
+    url,
+    path=None,
+    overwrite=True,
+    sha1_hash=None,
+    retries=5,
+    verify_ssl=True,
+    log=True,
+):
+    """Download a given URL.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    url : str
+        URL to download.
+    path : str, optional
+        Destination path to store downloaded file. By default stores to the
+        current directory with the same name as in url.
+    overwrite : bool, optional
+        Whether to overwrite the destination file if it already exists.
+        By default always overwrites the downloaded file.
+    sha1_hash : str, optional
+        Expected sha1 hash in hexadecimal digits. Will ignore existing file when hash is specified
+        but doesn't match.
+    retries : integer, default 5
+        The number of times to attempt downloading in case of failure or non 200 return codes.
+    verify_ssl : bool, default True
+        Verify SSL certificates.
+    log : bool, default True
+        Whether to print the progress for download
+
+    Returns
+    -------
+    str
+        The file path of the downloaded file.
+    """
+    if path is None:
+        fname = url.split("/")[-1]
+        # Empty filenames are invalid
+        assert fname, (
+            "Can't construct file-name from this URL. "
+            "Please set the `path` option manually."
+        )
+    else:
+        path = os.path.expanduser(path)
+        if os.path.isdir(path):
+            fname = os.path.join(path, url.split("/")[-1])
+        else:
+            fname = path
+    assert retries >= 0, "Number of retries should be at least 0"
+
+    if not verify_ssl:
+        warnings.warn(
+            "Unverified HTTPS request is being made (verify_ssl=False). "
+            "Adding certificate verification is strongly advised."
+        )
+
+    if (
+        overwrite
+        or not os.path.exists(fname)
+        or (sha1_hash and not check_sha1(fname, sha1_hash))
+    ):
+        dirname = os.path.dirname(os.path.abspath(os.path.expanduser(fname)))
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        while retries + 1 > 0:
+            # Disable pyling too broad Exception
+            # pylint: disable=W0703
+            try:
+                if log:
+                    print("Downloading %s from %s..." % (fname, url))
+                r = requests.get(url, stream=True, verify=verify_ssl)
+                if r.status_code != 200:
+                    raise RuntimeError("Failed downloading url %s" % url)
+                # Get the total file size.
+                total_size = int(r.headers.get("content-length", 0))
+                with tqdm(
+                    total=total_size, unit="B", unit_scale=True, desc=fname
+                ) as progress_bar:
+                    with open(fname, "wb") as f:
+                        for chunk in r.iter_content(chunk_size=1024):
+                            if chunk:  # filter out keep-alive new chunks
+                                f.write(chunk)
+                                progress_bar.update(len(chunk))
+                if sha1_hash and not check_sha1(fname, sha1_hash):
+                    raise UserWarning(
+                        "File {} is downloaded but the content hash does not match."
+                        " The repo may be outdated or download may be incomplete. "
+                        'If the "repo_url" is overridden, consider switching to '
+                        "the default repo.".format(fname)
+                    )
+                break
+            except Exception as e:
+                retries -= 1
+                if retries <= 0:
+                    raise e
+                if log:
+                    print(
+                        "download failed, retrying, {} attempt{} left".format(
+                            retries, "s" if retries > 1 else ""
+                        )
+                    )
+
+    return fname
+
+
+def check_sha1(filename, sha1_hash):
+    """Check whether the sha1 hash of the file content matches the expected hash.
+
+    Codes borrowed from mxnet/gluon/utils.py
+
+    Parameters
+    ----------
+    filename : str
+        Path to the file.
+    sha1_hash : str
+        Expected sha1 hash in hexadecimal digits.
+
+    Returns
+    -------
+    bool
+        Whether the file content matches the expected hash.
+    """
+    sha1 = hashlib.sha1()
+    with open(filename, "rb") as f:
+        while True:
+            data = f.read(1048576)
+            if not data:
+                break
+            sha1.update(data)
+
+    return sha1.hexdigest() == sha1_hash
+
+
+def extract_archive(file, target_dir, overwrite=True):
+    """Extract archive file.
+
+    Parameters
+    ----------
+    file : str
+        Absolute path of the archive file.
+    target_dir : str
+        Target directory of the archive to be uncompressed.
+    overwrite : bool, default True
+        Whether to overwrite the contents inside the directory.
+        By default always overwrites.
+    """
+    if os.path.exists(target_dir) and not overwrite:
+        return
+    print("Extracting file to {}".format(target_dir))
+    if (
+        file.endswith(".tar.gz")
+        or file.endswith(".tar")
+        or file.endswith(".tgz")
+    ):
+        import tarfile
+
+        with tarfile.open(file, "r") as archive:
+
+            def is_within_directory(directory, target):
+                abs_directory = os.path.abspath(directory)
+                abs_target = os.path.abspath(target)
+                prefix = os.path.commonprefix([abs_directory, abs_target])
+                return prefix == abs_directory
+
+            def safe_extract(
+                tar, path=".", members=None, *, numeric_owner=False
+            ):
+                for member in tar.getmembers():
+                    member_path = os.path.join(path, member.name)
+                    if not is_within_directory(path, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+                tar.extractall(path, members, numeric_owner=numeric_owner)
+
+            safe_extract(archive, path=target_dir)
+    elif file.endswith(".gz"):
+        import gzip
+        import shutil
+
+        with gzip.open(file, "rb") as f_in:
+            target_file = os.path.join(target_dir, os.path.basename(file)[:-3])
+            with open(target_file, "wb") as f_out:
+                shutil.copyfileobj(f_in, f_out)
+    elif file.endswith(".zip"):
+        import zipfile
+
+        with zipfile.ZipFile(file, "r") as archive:
+            archive.extractall(path=target_dir)
+    else:
+        raise Exception("Unrecognized file type: " + file)

python/dgl/graphbolt/item_sampler.py

@@ -0,0 +1,636 @@
+"""Item Sampler"""
+
+from collections.abc import Mapping
+from typing import Callable, Iterator, Optional, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import IterDataPipe
+
+from .internal import calculate_range
+from .internal_utils import gb_warning
+from .itemset import HeteroItemSet, ItemSet
+from .minibatch import MiniBatch
+
+__all__ = ["ItemSampler", "DistributedItemSampler", "minibatcher_default"]
+
+
+def minibatcher_default(batch, names):
+    """Default minibatcher which maps a list of items to a `MiniBatch` with the
+    same names as the items. The names of items are supposed to be provided
+    and align with the data attributes of `MiniBatch`. If any unknown item name
+    is provided, exception will be raised. If the names of items are not
+    provided, the item list is returned as is and a warning will be raised.
+
+    Parameters
+    ----------
+    batch : list
+        List of items.
+    names : Tuple[str] or None
+        Names of items in `batch` with same length. The order should align
+        with `batch`.
+
+    Returns
+    -------
+    MiniBatch
+        A minibatch.
+    """
+    if names is None:
+        gb_warning(
+            "Failed to map item list to `MiniBatch` as the names of items are "
+            "not provided. Please provide a customized `MiniBatcher`. "
+            "The item list is returned as is."
+        )
+        return batch
+    if len(names) == 1:
+        # Handle the case of single item: batch = tensor([0, 1, 2, 3]), names =
+        # ("seeds",) as `zip(batch, names)` will iterate over the tensor
+        # instead of the batch.
+        init_data = {names[0]: batch}
+    else:
+        if isinstance(batch, Mapping):
+            init_data = {
+                name: {k: v[i] for k, v in batch.items()}
+                for i, name in enumerate(names)
+            }
+        else:
+            init_data = {name: item for item, name in zip(batch, names)}
+    minibatch = MiniBatch()
+    # TODO(#7254): Hacks for original `seed_nodes` and `node_pairs`, which need
+    # to be cleaned up later.
+    if "node_pairs" in names:
+        pos_seeds = init_data["node_pairs"]
+        # Build negative graph.
+        if "negative_srcs" in names and "negative_dsts" in names:
+            neg_srcs = init_data["negative_srcs"]
+            neg_dsts = init_data["negative_dsts"]
+            (
+                init_data["seeds"],
+                init_data["labels"],
+                init_data["indexes"],
+            ) = _construct_seeds(
+                pos_seeds, neg_srcs=neg_srcs, neg_dsts=neg_dsts
+            )
+        elif "negative_srcs" in names:
+            neg_srcs = init_data["negative_srcs"]
+            (
+                init_data["seeds"],
+                init_data["labels"],
+                init_data["indexes"],
+            ) = _construct_seeds(pos_seeds, neg_srcs=neg_srcs)
+        elif "negative_dsts" in names:
+            neg_dsts = init_data["negative_dsts"]
+            (
+                init_data["seeds"],
+                init_data["labels"],
+                init_data["indexes"],
+            ) = _construct_seeds(pos_seeds, neg_dsts=neg_dsts)
+        else:
+            init_data["seeds"] = pos_seeds
+    for name, item in init_data.items():
+        if not hasattr(minibatch, name):
+            gb_warning(
+                f"Unknown item name '{name}' is detected and added into "
+                "`MiniBatch`. You probably need to provide a customized "
+                "`MiniBatcher`."
+            )
+        # TODO(#7254): Hacks for original `seed_nodes` and `node_pairs`, which
+        # need to be cleaned up later.
+        if name == "seed_nodes":
+            name = "seeds"
+        if name in ("node_pairs", "negative_srcs", "negative_dsts"):
+            continue
+        setattr(minibatch, name, item)
+    return minibatch
+
+
+class ItemSampler(IterDataPipe):
+    """A sampler to iterate over input items and create minibatches.
+
+    Input items could be node IDs, node pairs with or without labels, node
+    pairs with negative sources/destinations.
+
+    Note: This class `ItemSampler` is not decorated with
+    `torch.utils.data.functional_datapipe` on purpose. This indicates it
+    does not support function-like call. But any iterable datapipes from
+    `torch.utils.data.datapipes` can be further appended.
+
+    Parameters
+    ----------
+    item_set : Union[ItemSet, HeteroItemSet]
+        Data to be sampled.
+    batch_size : int
+        The size of each batch.
+    minibatcher : Optional[Callable]
+        A callable that takes in a list of items and returns a `MiniBatch`.
+    drop_last : bool
+        Option to drop the last batch if it's not full.
+    shuffle : bool
+        Option to shuffle before sample.
+    seed: int
+        The seed for reproducible stochastic shuffling. If None, a random seed
+        will be generated.
+
+    Examples
+    --------
+    1. Node IDs.
+
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> item_set = gb.ItemSet(torch.arange(0, 10), names="seeds")
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4, shuffle=False, drop_last=False
+    ... )
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds=tensor([0, 1, 2, 3]), sampled_subgraphs=None,
+        node_features=None, labels=None, input_nodes=None,
+        indexes=None, edge_features=None, compacted_seeds=None,
+        blocks=None,)
+
+    2. Node pairs.
+
+    >>> item_set = gb.ItemSet(torch.arange(0, 20).reshape(-1, 2),
+    ...     names="seeds")
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4, shuffle=False, drop_last=False
+    ... )
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds=tensor([[0, 1], [2, 3], [4, 5], [6, 7]]),
+        sampled_subgraphs=None, node_features=None, labels=None,
+        input_nodes=None, indexes=None, edge_features=None,
+        compacted_seeds=None, blocks=None,)
+
+    3. Node pairs and labels.
+
+    >>> item_set = gb.ItemSet(
+    ...     (torch.arange(0, 20).reshape(-1, 2), torch.arange(10, 20)),
+    ...     names=("seeds", "labels")
+    ... )
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4, shuffle=False, drop_last=False
+    ... )
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds=tensor([[0, 1], [2, 3], [4, 5], [6, 7]]),
+        sampled_subgraphs=None, node_features=None,
+        labels=tensor([10, 11, 12, 13]), input_nodes=None,
+        indexes=None, edge_features=None, compacted_seeds=None,
+        blocks=None,)
+
+    4. Node pairs, labels and indexes.
+
+    >>> seeds = torch.arange(0, 20).reshape(-1, 2)
+    >>> labels = torch.tensor([1, 1, 0, 0, 0, 0, 0, 0, 0, 0])
+    >>> indexes = torch.tensor([0, 1, 0, 0, 0, 0, 1, 1, 1, 1])
+    >>> item_set = gb.ItemSet((seeds, labels, indexes), names=("seeds",
+    ...     "labels", "indexes"))
+    >>> item_sampler = gb.ItemSampler(
+    ...     item_set, batch_size=4, shuffle=False, drop_last=False
+    ... )
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds=tensor([[0, 1], [2, 3], [4, 5], [6, 7]]),
+        sampled_subgraphs=None, node_features=None,
+        labels=tensor([1, 1, 0, 0]), input_nodes=None,
+        indexes=tensor([0, 1, 0, 0]), edge_features=None,
+        compacted_seeds=None, blocks=None,)
+
+    5. Further process batches with other datapipes such as
+    :class:`torch.utils.data.datapipes.iter.Mapper`.
+
+    >>> item_set = gb.ItemSet(torch.arange(0, 10))
+    >>> data_pipe = gb.ItemSampler(item_set, 4)
+    >>> def add_one(batch):
+    ...     return batch + 1
+    >>> data_pipe = data_pipe.map(add_one)
+    >>> list(data_pipe)
+    [tensor([1, 2, 3, 4]), tensor([5, 6, 7, 8]), tensor([ 9, 10])]
+
+    6. Heterogeneous node IDs.
+
+    >>> ids = {
+    ...     "user": gb.ItemSet(torch.arange(0, 5), names="seeds"),
+    ...     "item": gb.ItemSet(torch.arange(0, 6), names="seeds"),
+    ... }
+    >>> item_set = gb.HeteroItemSet(ids)
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=4)
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds={'user': tensor([0, 1, 2, 3])}, sampled_subgraphs=None,
+        node_features=None, labels=None, input_nodes=None, indexes=None,
+        edge_features=None, compacted_seeds=None, blocks=None,)
+
+    7. Heterogeneous node pairs.
+
+    >>> seeds_like = torch.arange(0, 10).reshape(-1, 2)
+    >>> seeds_follow = torch.arange(10, 20).reshape(-1, 2)
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user:like:item": gb.ItemSet(
+    ...         seeds_like, names="seeds"),
+    ...     "user:follow:user": gb.ItemSet(
+    ...         seeds_follow, names="seeds"),
+    ... })
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=4)
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds={'user:like:item':
+        tensor([[0, 1], [2, 3], [4, 5], [6, 7]])}, sampled_subgraphs=None,
+        node_features=None, labels=None, input_nodes=None, indexes=None,
+        edge_features=None, compacted_seeds=None, blocks=None,)
+
+    8. Heterogeneous node pairs and labels.
+
+    >>> seeds_like = torch.arange(0, 10).reshape(-1, 2)
+    >>> labels_like = torch.arange(0, 5)
+    >>> seeds_follow = torch.arange(10, 20).reshape(-1, 2)
+    >>> labels_follow = torch.arange(5, 10)
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user:like:item": gb.ItemSet((seeds_like, labels_like),
+    ...         names=("seeds", "labels")),
+    ...     "user:follow:user": gb.ItemSet((seeds_follow, labels_follow),
+    ...         names=("seeds", "labels")),
+    ... })
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=4)
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds={'user:like:item':
+        tensor([[0, 1], [2, 3], [4, 5], [6, 7]])}, sampled_subgraphs=None,
+        node_features=None, labels={'user:like:item': tensor([0, 1, 2, 3])},
+        input_nodes=None, indexes=None, edge_features=None,
+        compacted_seeds=None, blocks=None,)
+
+    9. Heterogeneous node pairs, labels and indexes.
+
+    >>> seeds_like = torch.arange(0, 10).reshape(-1, 2)
+    >>> labels_like = torch.tensor([1, 1, 0, 0, 0])
+    >>> indexes_like = torch.tensor([0, 1, 0, 0, 1])
+    >>> seeds_follow = torch.arange(20, 30).reshape(-1, 2)
+    >>> labels_follow = torch.tensor([1, 1, 0, 0, 0])
+    >>> indexes_follow = torch.tensor([0, 1, 0, 0, 1])
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user:like:item": gb.ItemSet((seeds_like, labels_like,
+    ...         indexes_like), names=("seeds", "labels", "indexes")),
+    ...     "user:follow:user": gb.ItemSet((seeds_follow,labels_follow,
+    ...         indexes_follow), names=("seeds", "labels", "indexes")),
+    ... })
+    >>> item_sampler = gb.ItemSampler(item_set, batch_size=4)
+    >>> next(iter(item_sampler))
+    MiniBatch(seeds={'user:like:item':
+        tensor([[0, 1], [2, 3], [4, 5], [6, 7]])}, sampled_subgraphs=None,
+        node_features=None, labels={'user:like:item': tensor([1, 1, 0, 0])},
+        input_nodes=None, indexes={'user:like:item': tensor([0, 1, 0, 0])},
+        edge_features=None, compacted_seeds=None, blocks=None,)
+    """
+
+    def __init__(
+        self,
+        item_set: Union[ItemSet, HeteroItemSet],
+        batch_size: int,
+        minibatcher: Optional[Callable] = minibatcher_default,
+        drop_last: Optional[bool] = False,
+        shuffle: Optional[bool] = False,
+        seed: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+        self._item_set = item_set
+        self._names = item_set.names
+        self._batch_size = batch_size
+        self._minibatcher = minibatcher
+        self._drop_last = drop_last
+        self._shuffle = shuffle
+        self._distributed = False
+        self._drop_uneven_inputs = False
+        self._world_size = None
+        self._rank = None
+        # For the sake of reproducibility, the seed should be allowed to be
+        # manually set by the user.
+        if seed is None:
+            self._seed = np.random.randint(0, np.iinfo(np.int32).max)
+        else:
+            self._seed = seed
+        # The attribute `self._epoch` is added to make shuffling work properly
+        # across multiple epochs. Otherwise, the same ordering will always be
+        # used in every epoch.
+        self._epoch = 0
+
+    def __iter__(self) -> Iterator:
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is not None:
+            num_workers = worker_info.num_workers
+            worker_id = worker_info.id
+        else:
+            num_workers = 1
+            worker_id = 0
+        total = len(self._item_set)
+        start_offset, assigned_count, output_count = calculate_range(
+            self._distributed,
+            total,
+            self._world_size,
+            self._rank,
+            num_workers,
+            worker_id,
+            self._batch_size,
+            self._drop_last,
+            self._drop_uneven_inputs,
+        )
+        if self._shuffle:
+            g = torch.Generator().manual_seed(self._seed + self._epoch)
+            permutation = torch.randperm(total, generator=g)
+            indices = permutation[start_offset : start_offset + assigned_count]
+        else:
+            indices = torch.arange(start_offset, start_offset + assigned_count)
+        for i in range(0, assigned_count, self._batch_size):
+            if output_count <= 0:
+                break
+            yield self._minibatcher(
+                self._item_set[
+                    indices[i : i + min(self._batch_size, output_count)]
+                ],
+                self._names,
+            )
+            output_count -= self._batch_size
+
+        self._epoch += 1
+
+
+class DistributedItemSampler(ItemSampler):
+    """A sampler to iterate over input items and create subsets distributedly.
+
+    This sampler creates a distributed subset of items from the given data set,
+    which can be used for training with PyTorch's Distributed Data Parallel
+    (DDP). The items can be node IDs, node pairs with or without labels, node
+    pairs with negative sources/destinations, DGLGraphs, or heterogeneous
+    counterparts. The original item set is split such that each replica
+    (process) receives an exclusive subset.
+
+    Note: The items will be first split onto each replica, then get shuffled
+    (if needed) and batched. Therefore, each replica will always get a same set
+    of items.
+
+    Note: This class `DistributedItemSampler` is not decorated with
+    `torch.utils.data.functional_datapipe` on purpose. This indicates it
+    does not support function-like call. But any iterable datapipes from
+    `torch.utils.data.datapipes` can be further appended.
+
+    Parameters
+    ----------
+    item_set : Union[ItemSet, HeteroItemSet]
+        Data to be sampled.
+    batch_size : int
+        The size of each batch.
+    minibatcher : Optional[Callable]
+        A callable that takes in a list of items and returns a `MiniBatch`.
+    drop_last : bool
+        Option to drop the last batch if it's not full.
+    shuffle : bool
+        Option to shuffle before sample.
+    num_replicas: int
+        The number of model replicas that will be created during Distributed
+        Data Parallel (DDP) training. It should be the same as the real world
+        size, otherwise it could cause errors. By default, it is retrieved from
+        the current distributed group.
+    drop_uneven_inputs : bool
+        Option to make sure the numbers of batches for each replica are the
+        same. If some of the replicas have more batches than the others, the
+        redundant batches of those replicas will be dropped. If the drop_last
+        parameter is also set to True, the last batch will be dropped before the
+        redundant batches are dropped.
+        Note: When using Distributed Data Parallel (DDP) training, the program
+        may hang or error if the a replica has fewer inputs. It is recommended
+        to use the Join Context Manager provided by PyTorch to solve this
+        problem. Please refer to
+        https://pytorch.org/tutorials/advanced/generic_join.html. However, this
+        option can be used if the Join Context Manager is not helpful for any
+        reason.
+    seed: int
+        The seed for reproducible stochastic shuffling. If None, a random seed
+        will be generated.
+
+    Examples
+    --------
+    0. Preparation: DistributedItemSampler needs multi-processing environment to
+    work. You need to spawn subprocesses and initialize processing group before
+    executing following examples. Due to randomness, the output is not always
+    the same as listed below.
+
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+    >>> item_set = gb.ItemSet(torch.arange(15))
+    >>> num_replicas = 4
+    >>> batch_size = 2
+    >>> mp.spawn(...)
+
+    1. shuffle = False, drop_last = False, drop_uneven_inputs = False.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=False, drop_last=False,
+    >>>     drop_uneven_inputs=False
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    Replica#0: [tensor([0, 1]), tensor([2, 3])]
+    Replica#1: [tensor([4, 5]), tensor([6, 7])]
+    Replica#2: [tensor([8, 9]), tensor([10, 11])]
+    Replica#3: [tensor([12, 13]), tensor([14])]
+
+    2. shuffle = False, drop_last = True, drop_uneven_inputs = False.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=False, drop_last=True,
+    >>>     drop_uneven_inputs=False
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    Replica#0: [tensor([0, 1]), tensor([2, 3])]
+    Replica#1: [tensor([4, 5]), tensor([6, 7])]
+    Replica#2: [tensor([8, 9]), tensor([10, 11])]
+    Replica#3: [tensor([12, 13])]
+
+    3. shuffle = False, drop_last = False, drop_uneven_inputs = True.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=False, drop_last=False,
+    >>>     drop_uneven_inputs=True
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    Replica#0: [tensor([0, 1]), tensor([2, 3])]
+    Replica#1: [tensor([4, 5]), tensor([6, 7])]
+    Replica#2: [tensor([8, 9]), tensor([10, 11])]
+    Replica#3: [tensor([12, 13]), tensor([14])]
+
+    4. shuffle = False, drop_last = True, drop_uneven_inputs = True.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=False, drop_last=True,
+    >>>     drop_uneven_inputs=True
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    Replica#0: [tensor([0, 1])]
+    Replica#1: [tensor([4, 5])]
+    Replica#2: [tensor([8, 9])]
+    Replica#3: [tensor([12, 13])]
+
+    5. shuffle = True, drop_last = True, drop_uneven_inputs = False.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=True, drop_last=True,
+    >>>     drop_uneven_inputs=False
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    (One possible output:)
+    Replica#0: [tensor([3, 2]), tensor([0, 1])]
+    Replica#1: [tensor([6, 5]), tensor([7, 4])]
+    Replica#2: [tensor([8, 10])]
+    Replica#3: [tensor([14, 12])]
+
+    6. shuffle = True, drop_last = True, drop_uneven_inputs = True.
+
+    >>> item_sampler = gb.DistributedItemSampler(
+    >>>     item_set, batch_size=2, shuffle=True, drop_last=True,
+    >>>     drop_uneven_inputs=True
+    >>> )
+    >>> data_loader = gb.DataLoader(item_sampler)
+    >>> print(f"Replica#{proc_id}: {list(data_loader)})
+    (One possible output:)
+    Replica#0: [tensor([1, 3])]
+    Replica#1: [tensor([7, 5])]
+    Replica#2: [tensor([11, 9])]
+    Replica#3: [tensor([13, 14])]
+    """
+
+    def __init__(
+        self,
+        item_set: Union[ItemSet, HeteroItemSet],
+        batch_size: int,
+        minibatcher: Optional[Callable] = minibatcher_default,
+        drop_last: Optional[bool] = False,
+        shuffle: Optional[bool] = False,
+        drop_uneven_inputs: Optional[bool] = False,
+        seed: Optional[int] = None,
+    ) -> None:
+        super().__init__(
+            item_set,
+            batch_size,
+            minibatcher,
+            drop_last,
+            shuffle,
+            seed,
+        )
+        self._distributed = True
+        self._drop_uneven_inputs = drop_uneven_inputs
+        if not dist.is_available():
+            raise RuntimeError(
+                "Distributed item sampler requires distributed package."
+            )
+        self._world_size = dist.get_world_size()
+        self._rank = dist.get_rank()
+        if self._world_size > 1:
+            # For the sake of reproducibility, the seed should be allowed to be
+            # manually set by the user.
+            self._align_seeds(src=0, seed=seed)
+
+    def _align_seeds(
+        self, src: Optional[int] = 0, seed: Optional[int] = None
+    ) -> None:
+        """Aligns seeds across distributed processes.
+
+        This method synchronizes seeds across distributed processes, ensuring
+        consistent randomness.
+
+        Parameters
+        ----------
+        src: int, optional
+            The source process rank. Defaults to 0.
+        seed: int, optional
+            The seed value to synchronize. If None, a random seed will be
+            generated. Defaults to None.
+        """
+        device = (
+            torch.cuda.current_device()
+            if torch.cuda.is_available() and dist.get_backend() == "nccl"
+            else "cpu"
+        )
+        if seed is None:
+            seed = np.random.randint(0, np.iinfo(np.int32).max)
+        if self._rank == src:
+            seed_tensor = torch.tensor(seed, dtype=torch.int32, device=device)
+        else:
+            seed_tensor = torch.empty([], dtype=torch.int32, device=device)
+        dist.broadcast(seed_tensor, src=src)
+        self._seed = seed_tensor.item()
+
+
+def _construct_seeds(pos_seeds, neg_srcs=None, neg_dsts=None):
+    # For homogeneous graph.
+    if isinstance(pos_seeds, torch.Tensor):
+        negative_ratio = neg_srcs.size(1) if neg_srcs else neg_dsts.size(1)
+        neg_srcs = (
+            neg_srcs
+            if neg_srcs is not None
+            else pos_seeds[:, 0].repeat_interleave(negative_ratio)
+        ).view(-1)
+        neg_dsts = (
+            neg_dsts
+            if neg_dsts is not None
+            else pos_seeds[:, 1].repeat_interleave(negative_ratio)
+        ).view(-1)
+        neg_seeds = torch.cat((neg_srcs, neg_dsts)).view(2, -1).T
+        seeds = torch.cat((pos_seeds, neg_seeds))
+        pos_seeds_num = pos_seeds.size(0)
+        labels = torch.empty(seeds.size(0), device=pos_seeds.device)
+        labels[:pos_seeds_num] = 1
+        labels[pos_seeds_num:] = 0
+        pos_indexes = torch.arange(
+            0,
+            pos_seeds_num,
+            device=pos_seeds.device,
+        )
+        neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
+        indexes = torch.cat((pos_indexes, neg_indexes))
+    # For heterogeneous graph.
+    else:
+        negative_ratio = (
+            list(neg_srcs.values())[0].size(1)
+            if neg_srcs
+            else list(neg_dsts.values())[0].size(1)
+        )
+        seeds = {}
+        labels = {}
+        indexes = {}
+        for etype in pos_seeds:
+            neg_src = (
+                neg_srcs[etype]
+                if neg_srcs is not None
+                else pos_seeds[etype][:, 0].repeat_interleave(negative_ratio)
+            ).view(-1)
+            neg_dst = (
+                neg_dsts[etype]
+                if neg_dsts is not None
+                else pos_seeds[etype][:, 1].repeat_interleave(negative_ratio)
+            ).view(-1)
+            seeds[etype] = torch.cat(
+                (
+                    pos_seeds[etype],
+                    torch.cat(
+                        (
+                            neg_src,
+                            neg_dst,
+                        )
+                    )
+                    .view(2, -1)
+                    .T,
+                )
+            )
+            pos_seeds_num = pos_seeds[etype].size(0)
+            labels[etype] = torch.empty(
+                seeds[etype].size(0), device=pos_seeds[etype].device
+            )
+            labels[etype][:pos_seeds_num] = 1
+            labels[etype][pos_seeds_num:] = 0
+            pos_indexes = torch.arange(
+                0,
+                pos_seeds_num,
+                device=pos_seeds[etype].device,
+            )
+            neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
+            indexes[etype] = torch.cat((pos_indexes, neg_indexes))
+    return seeds, labels, indexes

python/dgl/graphbolt/itemset.py

@@ -0,0 +1,454 @@
+"""GraphBolt Itemset."""
+
+import textwrap
+from typing import Dict, Iterable, Tuple, Union
+
+import torch
+
+from .internal_utils import gb_warning
+
+__all__ = ["ItemSet", "HeteroItemSet", "ItemSetDict"]
+
+
+def is_scalar(x):
+    """Checks if the input is a scalar."""
+    return (
+        len(x.shape) == 0 if isinstance(x, torch.Tensor) else isinstance(x, int)
+    )
+
+
+class ItemSet:
+    r"""A wrapper of a tensor or tuple of tensors.
+
+    Parameters
+    ----------
+    items: Union[int, torch.Tensor, Tuple[torch.Tensor]]
+        The tensors to be wrapped.
+        - If it is a single scalar (an integer or a tensor that holds a single
+          value), the item would be considered as a range_tensor created by
+          `torch.arange`.
+        - If it is a multi-dimensional tensor, the indexing will be performed
+          along the first dimension.
+        - If it is a tuple, each item in the tuple must be a tensor.
+
+    names: Union[str, Tuple[str]], optional
+        The names of the items. If it is a tuple, each name must corresponds to
+        an item in the `items` parameter. The naming is arbitrary, but in
+        general practice, the names should be chosen from ['labels', 'seeds',
+        'indexes'] to align with the attributes of class
+        `dgl.graphbolt.MiniBatch`.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+
+    1. Integer: number of nodes.
+
+    >>> num = 10
+    >>> item_set = gb.ItemSet(num, names="seeds")
+    >>> list(item_set)
+    [tensor(0), tensor(1), tensor(2), tensor(3), tensor(4), tensor(5),
+     tensor(6), tensor(7), tensor(8), tensor(9)]
+    >>> item_set[:]
+    tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+    >>> item_set.names
+    ('seeds',)
+
+    2. Torch scalar: number of nodes. Customizable dtype compared to Integer.
+
+    >>> num = torch.tensor(10, dtype=torch.int32)
+    >>> item_set = gb.ItemSet(num, names="seeds")
+    >>> list(item_set)
+    [tensor(0, dtype=torch.int32), tensor(1, dtype=torch.int32),
+     tensor(2, dtype=torch.int32), tensor(3, dtype=torch.int32),
+     tensor(4, dtype=torch.int32), tensor(5, dtype=torch.int32),
+     tensor(6, dtype=torch.int32), tensor(7, dtype=torch.int32),
+     tensor(8, dtype=torch.int32), tensor(9, dtype=torch.int32)]
+    >>> item_set[:]
+    tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int32)
+    >>> item_set.names
+    ('seeds',)
+
+    3. Single tensor: seed nodes.
+
+    >>> node_ids = torch.arange(0, 5)
+    >>> item_set = gb.ItemSet(node_ids, names="seeds")
+    >>> list(item_set)
+    [tensor(0), tensor(1), tensor(2), tensor(3), tensor(4)]
+    >>> item_set[:]
+    tensor([0, 1, 2, 3, 4])
+    >>> item_set.names
+    ('seeds',)
+
+    4. Tuple of tensors with same shape: seed nodes and labels.
+
+    >>> node_ids = torch.arange(0, 5)
+    >>> labels = torch.arange(5, 10)
+    >>> item_set = gb.ItemSet(
+    ...     (node_ids, labels), names=("seeds", "labels"))
+    >>> list(item_set)
+    [(tensor(0), tensor(5)), (tensor(1), tensor(6)), (tensor(2), tensor(7)),
+     (tensor(3), tensor(8)), (tensor(4), tensor(9))]
+    >>> item_set[:]
+    (tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9]))
+    >>> item_set.names
+    ('seeds', 'labels')
+
+    5. Tuple of tensors with different shape: seeds and labels.
+
+    >>> seeds = torch.arange(0, 10).reshape(-1, 2)
+    >>> labels = torch.tensor([1, 1, 0, 0, 0])
+    >>> item_set = gb.ItemSet(
+    ...     (seeds, labels), names=("seeds", "lables"))
+    >>> list(item_set)
+    [(tensor([0, 1]), tensor([1])),
+     (tensor([2, 3]), tensor([1])),
+     (tensor([4, 5]), tensor([0])),
+     (tensor([6, 7]), tensor([0])),
+     (tensor([8, 9]), tensor([0]))]
+    >>> item_set[:]
+    (tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
+     tensor([1, 1, 0, 0, 0]))
+    >>> item_set.names
+    ('seeds', 'labels')
+
+    6. Tuple of tensors with different shape: hyperlink and labels.
+
+    >>> seeds = torch.arange(0, 10).reshape(-1, 5)
+    >>> labels = torch.tensor([1, 0])
+    >>> item_set = gb.ItemSet(
+    ...     (seeds, labels), names=("seeds", "lables"))
+    >>> list(item_set)
+    [(tensor([0, 1, 2, 3, 4]), tensor([1])),
+     (tensor([5, 6, 7, 8, 9]), tensor([0]))]
+    >>> item_set[:]
+    (tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
+     tensor([1, 0]))
+    >>> item_set.names
+    ('seeds', 'labels')
+    """
+
+    def __init__(
+        self,
+        items: Union[int, torch.Tensor, Tuple[torch.Tensor]],
+        names: Union[str, Tuple[str]] = None,
+    ) -> None:
+        if is_scalar(items):
+            self._length = int(items)
+            self._items = items
+        elif isinstance(items, tuple):
+            self._length = len(items[0])
+            if any(self._length != len(item) for item in items):
+                raise ValueError("Size mismatch between items.")
+            self._items = items
+        else:
+            self._length = len(items)
+            self._items = (items,)
+        self._num_items = (
+            len(self._items) if isinstance(self._items, tuple) else 1
+        )
+        if names is not None:
+            if isinstance(names, tuple):
+                self._names = names
+            else:
+                self._names = (names,)
+            assert self._num_items == len(self._names), (
+                f"Number of items ({self._num_items}) and "
+                f"names ({len(self._names)}) don't match."
+            )
+        else:
+            self._names = None
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __getitem__(self, index: Union[int, slice, Iterable[int]]):
+        if is_scalar(self._items):
+            dtype = getattr(self._items, "dtype", torch.int64)
+            if isinstance(index, slice):
+                start, stop, step = index.indices(self._length)
+                return torch.arange(start, stop, step, dtype=dtype)
+            elif isinstance(index, int):
+                if index < 0:
+                    index += self._length
+                if index < 0 or index >= self._length:
+                    raise IndexError(
+                        f"{type(self).__name__} index out of range."
+                    )
+                return torch.tensor(index, dtype=dtype)
+            elif isinstance(index, torch.Tensor):
+                return index.to(dtype)
+            else:
+                raise TypeError(
+                    f"{type(self).__name__} indices must be int, slice, or "
+                    f"torch.Tensor, not {type(index)}."
+                )
+        elif self._num_items == 1:
+            return self._items[0][index]
+        else:
+            return tuple(item[index] for item in self._items)
+
+    @property
+    def names(self) -> Tuple[str]:
+        """Return the names of the items."""
+        return self._names
+
+    @property
+    def num_items(self) -> int:
+        """Return the number of the items."""
+        return self._num_items
+
+    def __repr__(self) -> str:
+        ret = (
+            f"{self.__class__.__name__}(\n"
+            f"    items={self._items},\n"
+            f"    names={self._names},\n"
+            f")"
+        )
+        return ret
+
+
+class HeteroItemSet:
+    r"""A collection of itemsets, each associated with a unique type.
+
+    This class aims to assemble existing itemsets with different types, for
+    example, seed_nodes of different node types in a graph.
+
+    Parameters
+    ----------
+    itemsets: Dict[str, ItemSet]
+        A dictionary whose keys are types and values are ItemSet instances.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from dgl import graphbolt as gb
+
+    1. Each itemset is a single tensor: seed nodes.
+
+    >>> node_ids_user = torch.arange(0, 5)
+    >>> node_ids_item = torch.arange(5, 10)
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user": gb.ItemSet(node_ids_user, names="seeds"),
+    ...     "item": gb.ItemSet(node_ids_item, names="seeds")})
+    >>> list(item_set)
+    [{"user": tensor(0)}, {"user": tensor(1)}, {"user": tensor(2)},
+     {"user": tensor(3)}, {"user": tensor(4)}, {"item": tensor(5)},
+     {"item": tensor(6)}, {"item": tensor(7)}, {"item": tensor(8)},
+     {"item": tensor(9)}}]
+    >>> item_set[:]
+    {"user": tensor([0, 1, 2, 3, 4]), "item": tensor([5, 6, 7, 8, 9])}
+    >>> item_set.names
+    ('seeds',)
+
+    2. Each itemset is a tuple of tensors with same shape: seed nodes and
+    labels.
+
+    >>> node_ids_user = torch.arange(0, 2)
+    >>> labels_user = torch.arange(0, 2)
+    >>> node_ids_item = torch.arange(2, 5)
+    >>> labels_item = torch.arange(2, 5)
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user": gb.ItemSet(
+    ...         (node_ids_user, labels_user),
+    ...         names=("seeds", "labels")),
+    ...     "item": gb.ItemSet(
+    ...         (node_ids_item, labels_item),
+    ...         names=("seeds", "labels"))})
+    >>> list(item_set)
+    [{"user": (tensor(0), tensor(0))}, {"user": (tensor(1), tensor(1))},
+     {"item": (tensor(2), tensor(2))}, {"item": (tensor(3), tensor(3))},
+     {"item": (tensor(4), tensor(4))}}]
+    >>> item_set[:]
+    {"user": (tensor([0, 1]), tensor([0, 1])),
+     "item": (tensor([2, 3, 4]), tensor([2, 3, 4]))}
+    >>> item_set.names
+    ('seeds', 'labels')
+
+    3. Each itemset is a tuple of tensors with different shape: seeds and
+    labels.
+
+    >>> seeds_like = torch.arange(0, 4).reshape(-1, 2)
+    >>> labels_like = torch.tensor([1, 0])
+    >>> seeds_follow = torch.arange(0, 6).reshape(-1, 2)
+    >>> labels_follow = torch.tensor([1, 1, 0])
+    >>> item_set = gb.HeteroItemSet({
+    ...     "user:like:item": gb.ItemSet(
+    ...         (seeds_like, labels_like),
+    ...         names=("seeds", "labels")),
+    ...     "user:follow:user": gb.ItemSet(
+    ...         (seeds_follow, labels_follow),
+    ...         names=("seeds", "labels"))})
+    >>> list(item_set)
+    [{'user:like:item': (tensor([0, 1]), tensor(1))},
+     {'user:like:item': (tensor([2, 3]), tensor(0))},
+     {'user:follow:user': (tensor([0, 1]), tensor(1))},
+     {'user:follow:user': (tensor([2, 3]), tensor(1))},
+     {'user:follow:user': (tensor([4, 5]), tensor(0))}]
+    >>> item_set[:]
+    {'user:like:item': (tensor([[0, 1], [2, 3]]),
+                        tensor([1, 0])),
+     'user:follow:user': (tensor([[0, 1], [2, 3], [4, 5]]),
+                          tensor([1, 1, 0]))}
+    >>> item_set.names
+    ('seeds', 'labels')
+
+    4. Each itemset is a tuple of tensors with different shape: hyperlink and
+    labels.
+
+    >>> first_seeds = torch.arange(0, 6).reshape(-1, 3)
+    >>> first_labels = torch.tensor([1, 0])
+    >>> second_seeds = torch.arange(0, 2).reshape(-1, 1)
+    >>> second_labels = torch.tensor([1, 0])
+    >>> item_set = gb.HeteroItemSet({
+    ...     "query:user:item": gb.ItemSet(
+    ...         (first_seeds, first_labels),
+    ...         names=("seeds", "labels")),
+    ...     "user": gb.ItemSet(
+    ...         (second_seeds, second_labels),
+    ...         names=("seeds", "labels"))})
+    >>> list(item_set)
+    [{'query:user:item': (tensor([0, 1, 2]), tensor(1))},
+     {'query:user:item': (tensor([3, 4, 5]), tensor(0))},
+     {'user': (tensor([0]), tensor(1))},
+     {'user': (tensor([1]), tensor(0))}]
+    >>> item_set[:]
+    {'query:user:item': (tensor([[0, 1, 2], [3, 4, 5]]),
+                        tensor([1, 0])),
+     'user': (tensor([[0], [1]]),tensor([1, 0]))}
+    >>> item_set.names
+    ('seeds', 'labels')
+    """
+
+    def __init__(self, itemsets: Dict[str, ItemSet]) -> None:
+        self._itemsets = itemsets
+        self._names = next(iter(itemsets.values())).names
+        assert all(
+            self._names == itemset.names for itemset in itemsets.values()
+        ), "All itemsets must have the same names."
+        offset = [0] + [len(itemset) for itemset in self._itemsets.values()]
+        self._offsets = torch.tensor(offset).cumsum(0)
+        self._length = int(self._offsets[-1])
+        self._keys = list(self._itemsets.keys())
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __getitem__(self, index: Union[int, slice, Iterable[int]]):
+        if isinstance(index, int):
+            if index < 0:
+                index += self._length
+            if index < 0 or index >= self._length:
+                raise IndexError(f"{type(self).__name__} index out of range.")
+            offset_idx = torch.searchsorted(self._offsets, index, right=True)
+            offset_idx -= 1
+            index -= self._offsets[offset_idx]
+            key = self._keys[offset_idx]
+            return {key: self._itemsets[key][index]}
+        elif isinstance(index, slice):
+            start, stop, step = index.indices(self._length)
+            if step != 1:
+                return self.__getitem__(torch.arange(start, stop, step))
+            assert start < stop, "Start must be smaller than stop."
+            data = {}
+            offset_idx_start = max(
+                1, torch.searchsorted(self._offsets, start, right=False)
+            )
+            for offset_idx in range(offset_idx_start, len(self._offsets)):
+                key = self._keys[offset_idx - 1]
+                data[key] = self._itemsets[key][
+                    max(0, start - self._offsets[offset_idx - 1]) : stop
+                    - self._offsets[offset_idx - 1]
+                ]
+                if stop <= self._offsets[offset_idx]:
+                    break
+            return data
+        elif isinstance(index, Iterable):
+            if not isinstance(index, torch.Tensor):
+                index = torch.tensor(index)
+            assert torch.all((index >= 0) & (index < self._length))
+            key_indices = (
+                torch.searchsorted(self._offsets, index, right=True) - 1
+            )
+            data = {}
+            for key_id, key in enumerate(self._keys):
+                mask = (key_indices == key_id).nonzero().squeeze(1)
+                if len(mask) == 0:
+                    continue
+                data[key] = self._itemsets[key][
+                    index[mask] - self._offsets[key_id]
+                ]
+            return data
+        else:
+            raise TypeError(
+                f"{type(self).__name__} indices must be int, slice, or "
+                f"iterable of int, not {type(index)}."
+            )
+
+    @property
+    def names(self) -> Tuple[str]:
+        """Return the names of the items."""
+        return self._names
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    itemsets={itemsets},\n"
+            "    names={names},\n"
+            ")"
+        )
+
+        itemsets_str = textwrap.indent(
+            repr(self._itemsets), " " * len("    itemsets=")
+        ).strip()
+
+        return ret.format(
+            Classname=self.__class__.__name__,
+            itemsets=itemsets_str,
+            names=self._names,
+        )
+
+
+class ItemSetDict:
+    """`ItemSetDict` is a deprecated class and will be removed in a future
+    version. Please use `HeteroItemSet` instead.
+
+    This class is an alias for `HeteroItemSet` and serves as a wrapper to
+    provide a smooth transition for users of the old class name. It issues a
+    deprecation warning upon instantiation and forwards all attribute access
+    and method calls to an instance of `HeteroItemSet`.
+    """
+
+    def __init__(self, itemsets: Dict[str, ItemSet]) -> None:
+        gb_warning(
+            "ItemSetDict is deprecated and will be removed in the future. "
+            "Please use HeteroItemSet instead.",
+            category=DeprecationWarning,
+        )
+        self._new_instance = HeteroItemSet(itemsets)
+
+    def __getattr__(self, name: str):
+        return getattr(self._new_instance, name)
+
+    def __getitem__(self, index):
+        return self._new_instance[index]
+
+    def __len__(self) -> int:
+        return len(self._new_instance)
+
+    def __repr__(self) -> str:
+        ret = (
+            "{Classname}(\n"
+            "    itemsets={itemsets},\n"
+            "    names={names},\n"
+            ")"
+        )
+        itemsets_str = textwrap.indent(
+            repr(self._itemsets), " " * len("    itemsets=")
+        ).strip()
+        return ret.format(
+            Classname=self.__class__.__name__,
+            itemsets=itemsets_str,
+            names=self._names,
+        )

python/dgl/graphbolt/minibatch.py

@@ -0,0 +1,417 @@
+"""Unified data structure for input and ouput of all the stages in loading process."""
+
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Union
+
+import torch
+
+from .base import (
+    apply_to,
+    CSCFormatBase,
+    etype_str_to_tuple,
+    expand_indptr,
+    is_object_pinned,
+)
+from .internal_utils import (
+    get_attributes,
+    get_nonproperty_attributes,
+    recursive_apply,
+)
+from .sampled_subgraph import SampledSubgraph
+
+__all__ = ["MiniBatch"]
+
+
+@dataclass
+class MiniBatch:
+    r"""A composite data class for data structure in the graphbolt.
+
+    It is designed to facilitate the exchange of data among different components
+    involved in processing data. The purpose of this class is to unify the
+    representation of input and output data across different stages, ensuring
+    consistency and ease of use throughout the loading process."""
+
+    labels: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Labels associated with seeds in the graph.
+    - If `labels` is a tensor: It indicates the graph is homogeneous. The value
+      should be corresponding labels to given 'seeds'.
+    - If `labels` is a dictionary: The keys should be node or edge type and the
+      value should be corresponding labels to given 'seeds'.
+    """
+
+    seeds: Union[
+        torch.Tensor,
+        Dict[str, torch.Tensor],
+    ] = None
+    """
+    Representation of seed items utilized in node classification tasks, link
+    prediction tasks and hyperlinks tasks.
+    - If `seeds` is a tensor: it indicates that the seeds originate from a
+      homogeneous graph. It can be either a 1-dimensional or 2-dimensional
+      tensor:
+        - 1-dimensional tensor: Each element directly represents a seed node
+          within the graph.
+        - 2-dimensional tensor: Each row designates a seed item, which can
+          encompass various entities such as edges, hyperlinks, or other graph
+          components depending on the specific context.
+    - If `seeds` is a dictionary: it indicates that the seeds originate from a
+      heterogeneous graph. The keys should be edge or node type, and the value
+      should be a tensor, which can be either a 1-dimensional or 2-dimensional
+      tensor:
+        - 1-dimensional tensor: Each element directly represents a seed node
+        of the given type within the graph.
+        - 2-dimensional tensor: Each row designates a seed item of the given
+          type, which can encompass various entities such as edges, hyperlinks,
+          or other graph components depending on the specific context.
+    """
+
+    indexes: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Indexes associated with seeds in the graph, which
+    indicates to which query a seeds belongs.
+    - If `indexes` is a tensor: It indicates the graph is homogeneous. The
+      value should be corresponding query to given 'seeds'.
+    - If `indexes` is a dictionary: It indicates the graph is heterogeneous.
+      The keys should be node or edge type and the value should be
+      corresponding query to given 'seeds'. For each key, indexes are
+      consecutive integers starting from zero.
+    """
+
+    sampled_subgraphs: List[SampledSubgraph] = None
+    """A list of 'SampledSubgraph's, each one corresponding to one layer,
+    representing a subset of a larger graph structure.
+    """
+
+    input_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """A representation of input nodes in the outermost layer. Conatins all nodes
+       in the 'sampled_subgraphs'.
+    - If `input_nodes` is a tensor: It indicates the graph is homogeneous.
+    - If `input_nodes` is a dictionary: The keys should be node type and the
+      value should be corresponding heterogeneous node id.
+    """
+
+    node_features: Union[
+        Dict[str, torch.Tensor], Dict[Tuple[str, str], torch.Tensor]
+    ] = None
+    """A representation of node features.
+      - If keys are single strings: It means the graph is homogeneous, and the
+      keys are feature names.
+      - If keys are tuples: It means the graph is heterogeneous, and the keys
+      are tuples of '(node_type, feature_name)'.
+    """
+
+    edge_features: List[
+        Union[Dict[str, torch.Tensor], Dict[Tuple[str, str], torch.Tensor]]
+    ] = None
+    """Edge features associated with the 'sampled_subgraphs'.
+      - If keys are single strings: It means the graph is homogeneous, and the
+      keys are feature names.
+      - If keys are tuples: It means the graph is heterogeneous, and the keys
+      are tuples of '(edge_type, feature_name)'. Note, edge type is single
+      string of format 'str:str:str'.
+    """
+
+    compacted_seeds: Union[
+        torch.Tensor,
+        Dict[str, torch.Tensor],
+    ] = None
+    """
+    Representation of compacted seeds corresponding to 'seeds', where
+    all node ids inside are compacted.
+    """
+
+    _blocks: list = None
+    """
+    A list of `DGLBlock`s.
+    """
+
+    def __repr__(self) -> str:
+        return _minibatch_str(self)
+
+    def node_ids(self) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        """A representation of input nodes in the outermost layer. Contains all
+        nodes in the `sampled_subgraphs`.
+        - If `input_nodes` is a tensor: It indicates the graph is homogeneous.
+        - If `input_nodes` is a dictionary: The keys should be node type and the
+          value should be corresponding heterogeneous node id.
+        """
+        return self.input_nodes
+
+    def num_layers(self) -> int:
+        """Return the number of layers."""
+        if self.sampled_subgraphs is None:
+            return 0
+        return len(self.sampled_subgraphs)
+
+    def edge_ids(
+        self, layer_id: int
+    ) -> Union[Dict[str, torch.Tensor], torch.Tensor]:
+        """Get the edge ids of a layer."""
+        return self.sampled_subgraphs[layer_id].original_edge_ids
+
+    def set_node_features(
+        self,
+        node_features: Union[
+            Dict[str, torch.Tensor], Dict[Tuple[str, str], torch.Tensor]
+        ],
+    ) -> None:
+        """Set node features."""
+        self.node_features = node_features
+
+    def set_edge_features(
+        self,
+        edge_features: List[
+            Union[Dict[str, torch.Tensor], Dict[Tuple[str, str], torch.Tensor]]
+        ],
+    ) -> None:
+        """Set edge features."""
+        self.edge_features = edge_features
+
+    @property
+    def blocks(self) -> list:
+        """DGL blocks extracted from `MiniBatch` containing graphical structures
+        and ID mappings.
+        """
+        if not self.sampled_subgraphs:
+            return None
+
+        if self._blocks is None:
+            self._blocks = self.compute_blocks()
+        return self._blocks
+
+    def compute_blocks(self) -> list:
+        """Extracts DGL blocks from `MiniBatch` to construct graphical
+        structures and ID mappings.
+        """
+        from dgl.convert import create_block, EID, NID
+
+        is_heterogeneous = isinstance(
+            self.sampled_subgraphs[0].sampled_csc, Dict
+        )
+
+        # Casts to minimum dtype in-place and returns self.
+        def cast_to_minimum_dtype(v: CSCFormatBase):
+            # Checks if number of vertices and edges fit into an int32.
+            dtype = (
+                torch.int32
+                if max(v.indptr.size(0) - 2, v.indices.size(0))
+                <= torch.iinfo(torch.int32).max
+                else torch.int64
+            )
+            v.indptr = v.indptr.to(dtype)
+            v.indices = v.indices.to(dtype)
+            return v
+
+        blocks = []
+        for subgraph in self.sampled_subgraphs:
+            original_row_node_ids = subgraph.original_row_node_ids
+            assert (
+                original_row_node_ids is not None
+            ), "Missing `original_row_node_ids` in sampled subgraph."
+            original_column_node_ids = subgraph.original_column_node_ids
+            assert (
+                original_column_node_ids is not None
+            ), "Missing `original_column_node_ids` in sampled subgraph."
+            if is_heterogeneous:
+                node_types = set()
+                sampled_csc = {}
+                for v in subgraph.sampled_csc.values():
+                    cast_to_minimum_dtype(v)
+                for etype, v in subgraph.sampled_csc.items():
+                    etype_tuple = etype_str_to_tuple(etype)
+                    node_types.add(etype_tuple[0])
+                    node_types.add(etype_tuple[2])
+                    sampled_csc[etype_tuple] = (
+                        "csc",
+                        (
+                            v.indptr,
+                            v.indices,
+                            torch.arange(
+                                0,
+                                len(v.indices),
+                                device=v.indptr.device,
+                                dtype=v.indptr.dtype,
+                            ),
+                        ),
+                    )
+                num_src_nodes = {
+                    ntype: (
+                        original_row_node_ids[ntype].size(0)
+                        if original_row_node_ids.get(ntype) is not None
+                        else 0
+                    )
+                    for ntype in node_types
+                }
+                num_dst_nodes = {
+                    ntype: (
+                        original_column_node_ids[ntype].size(0)
+                        if original_column_node_ids.get(ntype) is not None
+                        else 0
+                    )
+                    for ntype in node_types
+                }
+            else:
+                sampled_csc = cast_to_minimum_dtype(subgraph.sampled_csc)
+                sampled_csc = (
+                    "csc",
+                    (
+                        sampled_csc.indptr,
+                        sampled_csc.indices,
+                        torch.arange(
+                            0,
+                            len(sampled_csc.indices),
+                            device=sampled_csc.indptr.device,
+                            dtype=sampled_csc.indptr.dtype,
+                        ),
+                    ),
+                )
+                num_src_nodes = original_row_node_ids.size(0)
+                num_dst_nodes = original_column_node_ids.size(0)
+            blocks.append(
+                create_block(
+                    sampled_csc,
+                    num_src_nodes=num_src_nodes,
+                    num_dst_nodes=num_dst_nodes,
+                    node_count_check=False,
+                )
+            )
+
+        if is_heterogeneous:
+            # Assign reverse node ids to the outermost layer's source nodes.
+            for node_type, reverse_ids in self.sampled_subgraphs[
+                0
+            ].original_row_node_ids.items():
+                blocks[0].srcnodes[node_type].data[NID] = reverse_ids
+            # Assign reverse edges ids.
+            for block, subgraph in zip(blocks, self.sampled_subgraphs):
+                if subgraph.original_edge_ids is not None:
+                    for (
+                        edge_type,
+                        reverse_ids,
+                    ) in subgraph.original_edge_ids.items():
+                        block.edges[etype_str_to_tuple(edge_type)].data[
+                            EID
+                        ] = reverse_ids
+        else:
+            blocks[0].srcdata[NID] = self.sampled_subgraphs[
+                0
+            ].original_row_node_ids
+            # Assign reverse edges ids.
+            for block, subgraph in zip(blocks, self.sampled_subgraphs):
+                if subgraph.original_edge_ids is not None:
+                    block.edata[EID] = subgraph.original_edge_ids
+        return blocks
+
+    def to_pyg_data(self):
+        """Construct a PyG Data from `MiniBatch`. This function only supports
+        node classification task on a homogeneous graph and the number of
+        features cannot be more than one.
+        """
+        from torch_geometric.data import Data
+
+        if self.sampled_subgraphs is None:
+            edge_index = None
+        else:
+            col_nodes = []
+            row_nodes = []
+            for subgraph in self.sampled_subgraphs:
+                if subgraph is None:
+                    continue
+                sampled_csc = subgraph.sampled_csc
+                indptr = sampled_csc.indptr
+                indices = sampled_csc.indices
+                expanded_indptr = expand_indptr(
+                    indptr, dtype=indices.dtype, output_size=len(indices)
+                )
+                col_nodes.append(expanded_indptr)
+                row_nodes.append(indices)
+            col_nodes = torch.cat(col_nodes)
+            row_nodes = torch.cat(row_nodes)
+            edge_index = torch.unique(
+                torch.stack((row_nodes, col_nodes)), dim=1
+            ).long()
+
+        if self.node_features is None:
+            node_features = None
+        else:
+            assert (
+                len(self.node_features) == 1
+            ), "`to_pyg_data` only supports single feature homogeneous graph."
+            node_features = next(iter(self.node_features.values()))
+
+        if self.seeds is not None:
+            if isinstance(self.seeds, Dict):
+                batch_size = len(next(iter(self.seeds.values())))
+            else:
+                batch_size = len(self.seeds)
+        else:
+            batch_size = None
+        pyg_data = Data(
+            x=node_features,
+            edge_index=edge_index,
+            y=self.labels,
+            batch_size=batch_size,
+            n_id=self.node_ids(),
+        )
+        return pyg_data
+
+    def to(
+        self, device: torch.device, non_blocking=False
+    ):  # pylint: disable=invalid-name
+        """Copy `MiniBatch` to the specified device using reflection."""
+
+        copy_fn = lambda x: apply_to(x, device, non_blocking=non_blocking)
+
+        transfer_attrs = get_nonproperty_attributes(self)
+
+        for attr in transfer_attrs:
+            # Only copy member variables.
+            setattr(self, attr, recursive_apply(getattr(self, attr), copy_fn))
+
+        return self
+
+    def pin_memory(self):
+        """Copy `MiniBatch` to the pinned memory using reflection."""
+
+        return self.to("pinned")
+
+    def is_pinned(self) -> bool:
+        """Check whether `SampledSubgraph` is pinned using reflection."""
+
+        return is_object_pinned(self)
+
+
+def _minibatch_str(minibatch: MiniBatch) -> str:
+    final_str = ""
+    # Get all attributes in the class except methods.
+    attributes = get_attributes(minibatch)
+    attributes.reverse()
+    # Insert key with its value into the string.
+    for name in attributes:
+        if name[0] == "_":
+            continue
+        val = getattr(minibatch, name)
+
+        def _add_indent(_str, indent):
+            lines = _str.split("\n")
+            lines = [lines[0]] + [
+                " " * (indent + 10) + line for line in lines[1:]
+            ]
+            return "\n".join(lines)
+
+        # Let the variables in the list occupy one line each, and adjust the
+        # indentation on top of the original if the original data output has
+        # line feeds.
+        if isinstance(val, list):
+            val = [str(val_str) for val_str in val]
+            val = "[" + ",\n".join(val) + "]"
+        elif isinstance(val, tuple):
+            val = [str(val_str) for val_str in val]
+            val = "(" + ",\n".join(val) + ")"
+        else:
+            val = str(val)
+        final_str = (
+            final_str + f"{name}={_add_indent(val, len(name)+1)},\n" + " " * 10
+        )
+    return "MiniBatch(" + final_str[:-3] + ")"

python/dgl/graphbolt/minibatch_transformer.py

@@ -0,0 +1,46 @@
+"""Mini-batch transformer"""
+
+from torch.utils.data import functional_datapipe
+
+from torch.utils.data.datapipes.iter import Mapper
+
+from .minibatch import MiniBatch
+
+__all__ = [
+    "MiniBatchTransformer",
+]
+
+
+@functional_datapipe("transform")
+class MiniBatchTransformer(Mapper):
+    """A mini-batch transformer used to manipulate mini-batch.
+
+    Functional name: :obj:`transform`.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    transformer:
+        The function applied to each minibatch which is responsible for
+        transforming the minibatch.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        transformer=None,
+    ):
+        super().__init__(datapipe, self._transformer)
+        self.transformer = transformer or self._identity
+
+    def _transformer(self, minibatch):
+        minibatch = self.transformer(minibatch)
+        assert isinstance(
+            minibatch, (MiniBatch,)
+        ), "The transformer output should be an instance of MiniBatch"
+        return minibatch
+
+    @staticmethod
+    def _identity(minibatch):
+        return minibatch

python/dgl/graphbolt/negative_sampler.py

@@ -0,0 +1,104 @@
+"""Negative samplers."""
+
+from _collections_abc import Mapping
+
+from torch.utils.data import functional_datapipe
+
+from .minibatch_transformer import MiniBatchTransformer
+
+__all__ = [
+    "NegativeSampler",
+]
+
+
+@functional_datapipe("sample_negative")
+class NegativeSampler(MiniBatchTransformer):
+    """
+    A negative sampler used to generate negative samples and return
+    a mix of positive and negative samples.
+
+    Functional name: :obj:`sample_negative`.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    negative_ratio : int
+        The proportion of negative samples to positive samples.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        negative_ratio,
+    ):
+        super().__init__(datapipe, self._sample)
+        assert negative_ratio > 0, "Negative_ratio should be positive Integer."
+        self.negative_ratio = negative_ratio
+
+    def _sample(self, minibatch):
+        """
+        Generate a mix of positive and negative samples. If `seeds` in
+        minibatch is not None, `labels` and `indexes` will be constructed
+        after negative sampling, based on corresponding seeds.
+
+        Parameters
+        ----------
+        minibatch : MiniBatch
+            An instance of 'MiniBatch' class requires the 'seeds' field. This
+            function is responsible for generating negative edges corresponding
+            to the positive edges defined by the 'seeds'.
+
+        Returns
+        -------
+        MiniBatch
+            An instance of 'MiniBatch' encompasses both positive and negative
+            samples.
+        """
+        seeds = minibatch.seeds
+        if isinstance(seeds, Mapping):
+            if minibatch.indexes is None:
+                minibatch.indexes = {}
+            if minibatch.labels is None:
+                minibatch.labels = {}
+            for etype, pos_pairs in seeds.items():
+                (
+                    minibatch.seeds[etype],
+                    minibatch.labels[etype],
+                    minibatch.indexes[etype],
+                ) = self._sample_with_etype(pos_pairs, etype)
+        else:
+            (
+                minibatch.seeds,
+                minibatch.labels,
+                minibatch.indexes,
+            ) = self._sample_with_etype(seeds)
+        return minibatch
+
+    def _sample_with_etype(self, seeds, etype=None):
+        """Generate negative pairs for a given etype form positive pairs
+        for a given etype. If `seeds` is a 2D tensor, which represents
+        `seeds` is used in minibatch, corresponding labels and indexes will be
+        constructed.
+
+        Parameters
+        ----------
+        seeds : Tensor, Tensor
+            A N*2 tensors that represent source-destination node pairs of
+            positive edges, where positive means the edge must exist in the
+            graph.
+        etype : str
+            Canonical edge type.
+
+        Returns
+        -------
+        Tensor
+            A collection of postive and negative node pairs.
+        Tensor
+            Corresponding labels. If label is True, corresponding edge is
+            positive. If label is False, corresponding edge is negative.
+        Tensor
+            Corresponding indexes, indicates to which query an edge belongs.
+
+        """
+        raise NotImplementedError

python/dgl/graphbolt/sampled_subgraph.py

@@ -0,0 +1,450 @@
+"""Graphbolt sampled subgraph."""
+
+# pylint: disable= invalid-name
+from typing import Dict, NamedTuple, Tuple, Union
+
+import torch
+
+from .base import (
+    apply_to,
+    CSCFormatBase,
+    etype_str_to_tuple,
+    expand_indptr,
+    is_object_pinned,
+    isin,
+)
+
+from .internal_utils import recursive_apply
+
+
+__all__ = ["SampledSubgraph"]
+
+
+class _ExcludeEdgesWaiter:
+    def __init__(self, sampled_subgraph, index):
+        self.sampled_subgraph = sampled_subgraph
+        self.index = index
+
+    def wait(self):
+        """Returns the stored value when invoked."""
+        sampled_subgraph = self.sampled_subgraph
+        index = self.index
+        # Ensure there is no memory leak.
+        self.sampled_subgraph = self.index = None
+
+        if isinstance(index, dict):
+            for k in list(index.keys()):
+                index[k] = index[k].wait()
+        else:
+            index = index.wait()
+
+        return type(sampled_subgraph)(*_slice_subgraph(sampled_subgraph, index))
+
+
+class PyGLayerData(NamedTuple):
+    """A named tuple class to represent homogenous inputs to a PyG model layer.
+    The fields are x (input features), edge_index and size
+    (source and destination sizes).
+    """
+
+    x: torch.Tensor
+    edge_index: torch.Tensor
+    size: Tuple[int, int]
+
+
+class PyGLayerHeteroData(NamedTuple):
+    """A named tuple class to represent heterogenous inputs to a PyG model
+    layer. The fields are x (input features), edge_index and size
+    (source and destination sizes), and all fields are dictionaries.
+    """
+
+    x: Dict[str, torch.Tensor]
+    edge_index: Dict[str, torch.Tensor]
+    size: Dict[str, Tuple[int, int]]
+
+
+class SampledSubgraph:
+    r"""An abstract class for sampled subgraph. In the context of a
+    heterogeneous graph, each field should be of `Dict` type. Otherwise,
+    for homogeneous graphs, each field should correspond to its respective
+    value type."""
+
+    @property
+    def sampled_csc(
+        self,
+    ) -> Union[CSCFormatBase, Dict[str, CSCFormatBase],]:
+        """Returns the node pairs representing edges in csc format.
+          - If `sampled_csc` is a CSCFormatBase: It should be in the csc
+            format. `indptr` stores the index in the data array where each
+            column starts. `indices` stores the row indices of the non-zero
+            elements.
+          - If `sampled_csc` is a dictionary: The keys should be edge type and
+            the values should be corresponding node pairs. The ids inside is
+            heterogeneous ids.
+
+        Examples
+        --------
+        1. Homogeneous graph.
+
+        >>> import dgl.graphbolt as gb
+        >>> import torch
+        >>> sampled_csc = gb.CSCFormatBase(
+        ...     indptr=torch.tensor([0, 1, 2, 3]),
+        ...     indices=torch.tensor([0, 1, 2]))
+        >>> print(sampled_csc)
+        CSCFormatBase(indptr=tensor([0, 1, 2, 3]),
+                    indices=tensor([0, 1, 2]),
+        )
+
+        2. Heterogeneous graph.
+
+        >>> sampled_csc = {"A:relation:B": gb.CSCFormatBase(
+        ...     indptr=torch.tensor([0, 1, 2, 3]),
+        ...     indices=torch.tensor([0, 1, 2]))}
+        >>> print(sampled_csc)
+        {'A:relation:B': CSCFormatBase(indptr=tensor([0, 1, 2, 3]),
+                    indices=tensor([0, 1, 2]),
+        )}
+        """
+        raise NotImplementedError
+
+    @property
+    def original_column_node_ids(
+        self,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Returns corresponding reverse column node ids the original graph.
+        Column's reverse node ids in the original graph. A graph structure
+        can be treated as a coordinated row and column pair, and this is
+        the mapped ids of the column.
+          - If `original_column_node_ids` is a tensor: It represents the
+            original node ids.
+          - If `original_column_node_ids` is a dictionary: The keys should be
+            node type and the values should be corresponding original
+            heterogeneous node ids.
+        If present, it means column IDs are compacted, and `sampled_csc`
+        column IDs match these compacted ones.
+        """
+        return None
+
+    @property
+    def original_row_node_ids(
+        self,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Returns corresponding reverse row node ids the original graph.
+        Row's reverse node ids in the original graph. A graph structure
+        can be treated as a coordinated row and column pair, and this is
+        the mapped ids of the row.
+          - If `original_row_node_ids` is a tensor: It represents the original
+            node ids.
+          - If `original_row_node_ids` is a dictionary: The keys should be node
+            type and the values should be corresponding original heterogeneous
+            node ids.
+        If present, it means row IDs are compacted, and `sampled_csc`
+        row IDs match these compacted ones."""
+        return None
+
+    @property
+    def original_edge_ids(self) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        """Returns corresponding reverse edge ids the original graph.
+        Reverse edge ids in the original graph. This is useful when edge
+        features are needed.
+          - If `original_edge_ids` is a tensor: It represents the original edge
+            ids.
+          - If `original_edge_ids` is a dictionary: The keys should be edge
+            type and the values should be corresponding original heterogeneous
+            edge ids.
+        """
+        return None
+
+    def exclude_edges(
+        self,
+        edges: Union[
+            Dict[str, torch.Tensor],
+            torch.Tensor,
+        ],
+        assume_num_node_within_int32: bool = True,
+        async_op: bool = False,
+    ):
+        r"""Exclude edges from the sampled subgraph.
+
+        This function can be used with sampled subgraphs, regardless of
+        whether they have compacted row/column nodes or not. If the original
+        subgraph has compacted row or column nodes, the corresponding row or
+        column nodes in the returned subgraph will also be compacted.
+
+        Parameters
+        ----------
+        self : SampledSubgraph
+            The sampled subgraph.
+        edges : Union[torch.Tensor, Dict[str, torch.Tensor]]
+            Edges to exclude. If sampled subgraph is homogeneous, then `edges`
+            should be a N*2 tensors representing the edges to exclude. If
+            sampled subgraph is heterogeneous, then `edges` should be a
+            dictionary of edge types and the corresponding edges to exclude.
+        assume_num_node_within_int32: bool
+            If True, assumes the value of node IDs in the provided `edges` fall
+            within the int32 range, which can significantly enhance computation
+            speed. Default: True
+        async_op: bool
+            Boolean indicating whether the call is asynchronous. If so, the
+            result can be obtained by calling wait on the returned future.
+
+        Returns
+        -------
+        SampledSubgraph
+            An instance of a class that inherits from `SampledSubgraph`.
+
+        Examples
+        --------
+        >>> import dgl.graphbolt as gb
+        >>> import torch
+        >>> sampled_csc = {"A:relation:B": gb.CSCFormatBase(
+        ...     indptr=torch.tensor([0, 1, 2, 3]),
+        ...     indices=torch.tensor([0, 1, 2]))}
+        >>> original_column_node_ids = {"B": torch.tensor([10, 11, 12])}
+        >>> original_row_node_ids = {"A": torch.tensor([13, 14, 15])}
+        >>> original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+        >>> subgraph = gb.SampledSubgraphImpl(
+        ...     sampled_csc=sampled_csc,
+        ...     original_column_node_ids=original_column_node_ids,
+        ...     original_row_node_ids=original_row_node_ids,
+        ...     original_edge_ids=original_edge_ids
+        ... )
+        >>> edges_to_exclude = {"A:relation:B": torch.tensor([[14, 11], [15, 12]])}
+        >>> result = subgraph.exclude_edges(edges_to_exclude)
+        >>> print(result.sampled_csc)
+        {'A:relation:B': CSCFormatBase(indptr=tensor([0, 1, 1, 1]),
+                    indices=tensor([0]),
+        )}
+        >>> print(result.original_column_node_ids)
+        {'B': tensor([10, 11, 12])}
+        >>> print(result.original_row_node_ids)
+        {'A': tensor([13, 14, 15])}
+        >>> print(result.original_edge_ids)
+        {'A:relation:B': tensor([19])}
+        """
+        # TODO: Add support for value > in32, then remove this line.
+        assert (
+            assume_num_node_within_int32
+        ), "Values > int32 are not supported yet."
+        assert (isinstance(self.sampled_csc, CSCFormatBase)) == isinstance(
+            edges, torch.Tensor
+        ), (
+            "The sampled subgraph and the edges to exclude should be both "
+            "homogeneous or both heterogeneous."
+        )
+        # Get type of calling class.
+        calling_class = type(self)
+
+        # Three steps to exclude edges:
+        # 1. Convert the node pairs to the original ids if they are compacted.
+        # 2. Exclude the edges and get the index of the edges to keep.
+        # 3. Slice the subgraph according to the index.
+        if isinstance(self.sampled_csc, CSCFormatBase):
+            reverse_edges = _to_reverse_ids(
+                self.sampled_csc,
+                self.original_row_node_ids,
+                self.original_column_node_ids,
+            )
+            index = _exclude_homo_edges(
+                reverse_edges, edges, assume_num_node_within_int32, async_op
+            )
+        else:
+            index = {}
+            for etype, pair in self.sampled_csc.items():
+                if etype not in edges:
+                    # No edges need to be excluded.
+                    index[etype] = None
+                    continue
+                src_type, _, dst_type = etype_str_to_tuple(etype)
+                original_row_node_ids = (
+                    None
+                    if self.original_row_node_ids is None
+                    else self.original_row_node_ids.get(src_type)
+                )
+                original_column_node_ids = (
+                    None
+                    if self.original_column_node_ids is None
+                    else self.original_column_node_ids.get(dst_type)
+                )
+                reverse_edges = _to_reverse_ids(
+                    pair,
+                    original_row_node_ids,
+                    original_column_node_ids,
+                )
+                index[etype] = _exclude_homo_edges(
+                    reverse_edges,
+                    edges[etype],
+                    assume_num_node_within_int32,
+                    async_op,
+                )
+        if async_op:
+            return _ExcludeEdgesWaiter(self, index)
+        else:
+            return calling_class(*_slice_subgraph(self, index))
+
+    def to_pyg(
+        self, x: Union[torch.Tensor, Dict[str, torch.Tensor]]
+    ) -> Union[PyGLayerData, PyGLayerHeteroData]:
+        """
+        Process layer inputs so that they can be consumed by a PyG model layer.
+
+        Parameters
+        ----------
+        x : Union[torch.Tensor, Dict[str, torch.Tensor]]
+            The input node features to the GNN layer.
+
+        Returns
+        -------
+        Union[PyGLayerData, PyGLayerHeteroData]
+            A named tuple class with `x`, `edge_index` and `size` fields.
+            Typically, a PyG GNN layer's forward method will accept these as
+            arguments.
+        """
+        if isinstance(x, torch.Tensor):
+            # Homogenous
+            src = self.sampled_csc.indices
+            dst = expand_indptr(
+                self.sampled_csc.indptr,
+                dtype=src.dtype,
+                output_size=src.size(0),
+            )
+            edge_index = torch.stack([src, dst], dim=0).long()
+            dst_size = self.sampled_csc.indptr.size(0) - 1
+            # h and h[:dst_size] correspond to source and destination features resp.
+            return PyGLayerData(
+                (x, x[:dst_size]), edge_index, (x.size(0), dst_size)
+            )
+        else:
+            # Heterogenous
+            x_dst_dict = {}
+            edge_index_dict = {}
+            sizes_dict = {}
+            for etype, sampled_csc in self.sampled_csc.items():
+                src = sampled_csc.indices
+                dst = expand_indptr(
+                    sampled_csc.indptr,
+                    dtype=src.dtype,
+                    output_size=src.size(0),
+                )
+                edge_index = torch.stack([src, dst], dim=0).long()
+                dst_size = sampled_csc.indptr.size(0) - 1
+                # h and h[:dst_size] correspond to source and destination features resp.
+                src_ntype, _, dst_ntype = etype_str_to_tuple(etype)
+                x_dst_dict[dst_ntype] = x[dst_ntype][:dst_size]
+                edge_index_dict[etype] = edge_index
+                sizes_dict[etype] = (x[src_ntype].size(0), dst_size)
+
+            return PyGLayerHeteroData(
+                (x, x_dst_dict), edge_index_dict, sizes_dict
+            )
+
+    def to(
+        self, device: torch.device, non_blocking=False
+    ) -> None:  # pylint: disable=invalid-name
+        """Copy `SampledSubgraph` to the specified device using reflection."""
+
+        for attr in dir(self):
+            # Only copy member variables.
+            if not callable(getattr(self, attr)) and not attr.startswith("__"):
+                setattr(
+                    self,
+                    attr,
+                    recursive_apply(
+                        getattr(self, attr),
+                        apply_to,
+                        device,
+                        non_blocking=non_blocking,
+                    ),
+                )
+
+        return self
+
+    def pin_memory(self):
+        """Copy `SampledSubgraph` to the pinned memory using reflection."""
+
+        return self.to("pinned")
+
+    def is_pinned(self) -> bool:
+        """Check whether `SampledSubgraph` is pinned using reflection."""
+
+        return is_object_pinned(self)
+
+
+def _to_reverse_ids(node_pair, original_row_node_ids, original_column_node_ids):
+    indptr = node_pair.indptr
+    indices = node_pair.indices
+    if original_row_node_ids is not None:
+        indices = torch.index_select(
+            original_row_node_ids, dim=0, index=indices
+        )
+    indptr = expand_indptr(
+        indptr, indices.dtype, original_column_node_ids, len(indices)
+    )
+    return (indices, indptr)
+
+
+def _relabel_two_arrays(lhs_array, rhs_array):
+    """Relabel two arrays into a consecutive range starting from 0."""
+    concated = torch.cat([lhs_array, rhs_array])
+    _, mapping = torch.unique(concated, return_inverse=True)
+    return mapping[: lhs_array.numel()], mapping[lhs_array.numel() :]
+
+
+def _exclude_homo_edges(
+    edges: Tuple[torch.Tensor, torch.Tensor],
+    edges_to_exclude: torch.Tensor,
+    assume_num_node_within_int32: bool,
+    async_op: bool,
+):
+    """Return the indices of edges to be included."""
+    if assume_num_node_within_int32:
+        val = edges[0].long() << 32 | edges[1].long()
+        edges_to_exclude_trans = edges_to_exclude.T
+        val_to_exclude = (
+            edges_to_exclude_trans[0].long() << 32
+            | edges_to_exclude_trans[1].long()
+        )
+    else:
+        # TODO: Add support for value > int32.
+        raise NotImplementedError(
+            "Values out of range int32 are not supported yet"
+        )
+    if async_op:
+        return torch.ops.graphbolt.is_not_in_index_async(val, val_to_exclude)
+    else:
+        mask = ~isin(val, val_to_exclude)
+        return torch.nonzero(mask, as_tuple=True)[0]
+
+
+def _slice_subgraph(subgraph: SampledSubgraph, index: torch.Tensor):
+    """Slice the subgraph according to the index."""
+
+    def _index_select(obj, index):
+        if obj is None:
+            return None
+        if index is None:
+            return obj
+        if isinstance(obj, CSCFormatBase):
+            new_indices = obj.indices[index]
+            new_indptr = torch.searchsorted(index, obj.indptr)
+            return CSCFormatBase(
+                indptr=new_indptr,
+                indices=new_indices,
+            )
+        if isinstance(obj, torch.Tensor):
+            return obj[index]
+        # Handle the case when obj is a dictionary.
+        assert isinstance(obj, dict)
+        assert isinstance(index, dict)
+        ret = {}
+        for k, v in obj.items():
+            ret[k] = _index_select(v, index[k])
+        return ret
+
+    return (
+        _index_select(subgraph.sampled_csc, index),
+        subgraph.original_column_node_ids,
+        subgraph.original_row_node_ids,
+        _index_select(subgraph.original_edge_ids, index),
+    )

python/dgl/graphbolt/sampling_graph.py

@@ -0,0 +1,86 @@
+"""Sampling Graphs."""
+
+from typing import Dict, Union
+
+import torch
+
+
+__all__ = ["SamplingGraph"]
+
+
+class SamplingGraph:
+    r"""Class for sampling graph."""
+
+    def __init__(self):
+        pass
+
+    def __repr__(self) -> str:
+        """Return a string representation of the graph.
+
+        Returns
+        -------
+        str
+            String representation of the graph.
+        """
+        raise NotImplementedError
+
+    @property
+    def num_nodes(self) -> Union[int, Dict[str, int]]:
+        """The number of nodes in the graph.
+        - If the graph is homogenous, returns an integer.
+        - If the graph is heterogenous, returns a dictionary.
+
+        Returns
+        -------
+        Union[int, Dict[str, int]]
+            The number of nodes. Integer indicates the total nodes number of a
+            homogenous graph; dict indicates nodes number per node types of a
+            heterogenous graph.
+        """
+        raise NotImplementedError
+
+    @property
+    def num_edges(self) -> Union[int, Dict[str, int]]:
+        """The number of edges in the graph.
+        - If the graph is homogenous, returns an integer.
+        - If the graph is heterogenous, returns a dictionary.
+
+        Returns
+        -------
+        Union[int, Dict[str, int]]
+            The number of edges. Integer indicates the total edges number of a
+            homogenous graph; dict indicates edges number per edge types of a
+            heterogenous graph.
+        """
+        raise NotImplementedError
+
+    def copy_to_shared_memory(self, shared_memory_name: str) -> "SamplingGraph":
+        """Copy the graph to shared memory.
+
+        Parameters
+        ----------
+        shared_memory_name : str
+            Name of the shared memory.
+
+        Returns
+        -------
+        SamplingGraph
+            The copied SamplingGraph object on shared memory.
+        """
+        raise NotImplementedError
+
+    # pylint: disable=invalid-name
+    def to(self, device: torch.device) -> "SamplingGraph":
+        """Copy graph to the specified device.
+
+        Parameters
+        ----------
+        device : torch.device
+            The destination device.
+
+        Returns
+        -------
+        SamplingGraph
+            The graph on the specified device.
+        """
+        raise NotImplementedError

python/dgl/graphbolt/subgraph_sampler.py

@@ -0,0 +1,591 @@
+"""Subgraph samplers"""
+
+from collections import defaultdict
+from functools import partial
+from typing import Dict
+
+import torch
+import torch.distributed as thd
+from torch.utils.data import functional_datapipe
+
+from .base import seed_type_str_to_ntypes
+from .internal import compact_temporal_nodes, unique_and_compact
+from .minibatch import MiniBatch
+from .minibatch_transformer import MiniBatchTransformer
+
+__all__ = [
+    "SubgraphSampler",
+    "all_to_all",
+    "convert_to_hetero",
+    "revert_to_homo",
+]
+
+
+class _NoOpWaiter:
+    def __init__(self, result):
+        self.result = result
+
+    def wait(self):
+        """Returns the stored value when invoked."""
+        result = self.result
+        # Ensure there is no memory leak.
+        self.result = None
+        return result
+
+
+def _shift(inputs: list, group=None):
+    cutoff = len(inputs) - thd.get_rank(group)
+    return inputs[cutoff:] + inputs[:cutoff]
+
+
+def all_to_all(outputs, inputs, group=None, async_op=False):
+    """Wrapper for thd.all_to_all that permuted outputs and inputs before
+    calling it. The arguments have the permutation
+    `rank, ..., world_size - 1, 0, ..., rank - 1` and we make it
+    `0, world_size - 1` before calling `thd.all_to_all`."""
+    shift_fn = partial(_shift, group=group)
+    outputs = shift_fn(list(outputs))
+    inputs = shift_fn(list(inputs))
+    if outputs[0].is_cuda:
+        return thd.all_to_all(outputs, inputs, group, async_op)
+    # gloo backend will be used.
+    outputs_single = torch.cat(outputs)
+    output_split_sizes = [o.size(0) for o in outputs]
+    handle = thd.all_to_all_single(
+        outputs_single,
+        torch.cat(inputs),
+        output_split_sizes,
+        [i.size(0) for i in inputs],
+        group,
+        async_op,
+    )
+    temp_outputs = outputs_single.split(output_split_sizes)
+
+    class _Waiter:
+        def __init__(self, handle, outputs, temp_outputs):
+            self.handle = handle
+            self.outputs = outputs
+            self.temp_outputs = temp_outputs
+
+        def wait(self):
+            """Returns the stored value when invoked."""
+            handle = self.handle
+            outputs = self.outputs
+            temp_outputs = self.temp_outputs
+            # Ensure that there is no leak
+            self.handle = self.outputs = self.temp_outputs = None
+
+            if handle is not None:
+                handle.wait()
+            for output, temp_output in zip(outputs, temp_outputs):
+                output.copy_(temp_output)
+
+    post_processor = _Waiter(handle, outputs, temp_outputs)
+    return post_processor if async_op else post_processor.wait()
+
+
+def revert_to_homo(d: dict):
+    """Utility function to convert a dictionary that stores homogenous data."""
+    is_homogenous = len(d) == 1 and "_N" in d
+    return list(d.values())[0] if is_homogenous else d
+
+
+def convert_to_hetero(item):
+    """Utility function to convert homogenous data to heterogenous with a single
+    node type."""
+    is_heterogenous = isinstance(item, dict)
+    return item if is_heterogenous else {"_N": item}
+
+
+@functional_datapipe("sample_subgraph")
+class SubgraphSampler(MiniBatchTransformer):
+    """A subgraph sampler used to sample a subgraph from a given set of nodes
+    from a larger graph.
+
+    Functional name: :obj:`sample_subgraph`.
+
+    This class is the base class of all subgraph samplers. Any subclass of
+    SubgraphSampler should implement either the :meth:`sample_subgraphs` method
+    or the :meth:`sampling_stages` method to define the fine-grained sampling
+    stages to take advantage of optimizations provided by the GraphBolt
+    DataLoader.
+
+    Parameters
+    ----------
+    datapipe : DataPipe
+        The datapipe.
+    args : Non-Keyword Arguments
+        Arguments to be passed into sampling_stages.
+    kwargs : Keyword Arguments
+        Arguments to be passed into sampling_stages. Preprocessing stage makes
+        use of the `asynchronous` and `cooperative` parameters before they are
+        passed to the sampling stages.
+    """
+
+    def __init__(
+        self,
+        datapipe,
+        *args,
+        **kwargs,
+    ):
+        async_op = kwargs.get("asynchronous", False)
+        cooperative = kwargs.get("cooperative", False)
+        preprocess_fn = partial(
+            self._preprocess, cooperative=cooperative, async_op=async_op
+        )
+        datapipe = datapipe.transform(preprocess_fn)
+        if async_op:
+            fn = partial(self._wait_preprocess_future, cooperative=cooperative)
+            datapipe = datapipe.buffer().transform(fn)
+        if cooperative:
+            datapipe = datapipe.transform(self._seeds_cooperative_exchange_1)
+            datapipe = datapipe.buffer()
+            datapipe = datapipe.transform(
+                self._seeds_cooperative_exchange_1_wait_future
+            ).buffer()
+            datapipe = datapipe.transform(self._seeds_cooperative_exchange_2)
+            datapipe = datapipe.buffer()
+            datapipe = datapipe.transform(self._seeds_cooperative_exchange_3)
+            datapipe = datapipe.buffer()
+            datapipe = datapipe.transform(self._seeds_cooperative_exchange_4)
+        datapipe = self.sampling_stages(datapipe, *args, **kwargs)
+        datapipe = datapipe.transform(self._postprocess)
+        super().__init__(datapipe)
+
+    @staticmethod
+    def _postprocess(minibatch):
+        delattr(minibatch, "_seed_nodes")
+        delattr(minibatch, "_seeds_timestamp")
+        return minibatch
+
+    @staticmethod
+    def _preprocess(minibatch, cooperative: bool, async_op: bool):
+        if minibatch.seeds is None:
+            raise ValueError(
+                f"Invalid minibatch {minibatch}: `seeds` should have a value."
+            )
+        rank = thd.get_rank() if cooperative else 0
+        world_size = thd.get_world_size() if cooperative else 1
+        results = SubgraphSampler._seeds_preprocess(
+            minibatch, rank, world_size, async_op
+        )
+        if async_op:
+            minibatch._preprocess_future = results
+        else:
+            (
+                minibatch._seed_nodes,
+                minibatch._seeds_timestamp,
+                minibatch.compacted_seeds,
+                offsets,
+            ) = results
+            if cooperative:
+                minibatch._seeds_offsets = offsets
+        return minibatch
+
+    @staticmethod
+    def _wait_preprocess_future(minibatch, cooperative: bool):
+        (
+            minibatch._seed_nodes,
+            minibatch._seeds_timestamp,
+            minibatch.compacted_seeds,
+            offsets,
+        ) = minibatch._preprocess_future.wait()
+        delattr(minibatch, "_preprocess_future")
+        if cooperative:
+            minibatch._seeds_offsets = offsets
+        return minibatch
+
+    @staticmethod
+    def _seeds_cooperative_exchange_1(minibatch):
+        rank = thd.get_rank()
+        world_size = thd.get_world_size()
+        seeds = minibatch._seed_nodes
+        is_homogeneous = not isinstance(seeds, dict)
+        if is_homogeneous:
+            seeds = {"_N": seeds}
+        if minibatch._seeds_offsets is None:
+            assert minibatch.compacted_seeds is None
+            minibatch._rank_sort_future = torch.ops.graphbolt.rank_sort_async(
+                list(seeds.values()), rank, world_size
+            )
+        return minibatch
+
+    @staticmethod
+    def _seeds_cooperative_exchange_1_wait_future(minibatch):
+        world_size = thd.get_world_size()
+        seeds = minibatch._seed_nodes
+        is_homogeneous = not isinstance(seeds, dict)
+        if is_homogeneous:
+            seeds = {"_N": seeds}
+        num_ntypes = len(seeds.keys())
+        if minibatch._seeds_offsets is None:
+            result = minibatch._rank_sort_future.wait()
+            delattr(minibatch, "_rank_sort_future")
+            sorted_seeds, sorted_compacted, sorted_offsets = {}, {}, {}
+            for i, (
+                seed_type,
+                (typed_sorted_seeds, typed_index, typed_offsets),
+            ) in enumerate(zip(seeds.keys(), result)):
+                sorted_seeds[seed_type] = typed_sorted_seeds
+                sorted_compacted[seed_type] = typed_index
+                sorted_offsets[seed_type] = typed_offsets
+
+            minibatch._seed_nodes = sorted_seeds
+            minibatch.compacted_seeds = revert_to_homo(sorted_compacted)
+            minibatch._seeds_offsets = sorted_offsets
+        else:
+            minibatch._seeds_offsets = {"_N": minibatch._seeds_offsets}
+        counts_sent = torch.empty(world_size * num_ntypes, dtype=torch.int64)
+        for i, offsets in enumerate(minibatch._seeds_offsets.values()):
+            counts_sent[
+                torch.arange(i, world_size * num_ntypes, num_ntypes)
+            ] = offsets.diff()
+        delattr(minibatch, "_seeds_offsets")
+        counts_received = torch.empty_like(counts_sent)
+        minibatch._counts_future = all_to_all(
+            counts_received.split(num_ntypes),
+            counts_sent.split(num_ntypes),
+            async_op=True,
+        )
+        minibatch._counts_sent = counts_sent
+        minibatch._counts_received = counts_received
+        return minibatch
+
+    @staticmethod
+    def _seeds_cooperative_exchange_2(minibatch):
+        world_size = thd.get_world_size()
+        seeds = minibatch._seed_nodes
+        minibatch._counts_future.wait()
+        delattr(minibatch, "_counts_future")
+        num_ntypes = len(seeds.keys())
+        seeds_received = {}
+        counts_sent = {}
+        counts_received = {}
+        for i, (ntype, typed_seeds) in enumerate(seeds.items()):
+            idx = torch.arange(i, world_size * num_ntypes, num_ntypes)
+            typed_counts_sent = minibatch._counts_sent[idx].tolist()
+            typed_counts_received = minibatch._counts_received[idx].tolist()
+            typed_seeds_received = typed_seeds.new_empty(
+                sum(typed_counts_received)
+            )
+            all_to_all(
+                typed_seeds_received.split(typed_counts_received),
+                typed_seeds.split(typed_counts_sent),
+            )
+            seeds_received[ntype] = typed_seeds_received
+            counts_sent[ntype] = typed_counts_sent
+            counts_received[ntype] = typed_counts_received
+        minibatch._seed_nodes = seeds_received
+        minibatch._counts_sent = revert_to_homo(counts_sent)
+        minibatch._counts_received = revert_to_homo(counts_received)
+        return minibatch
+
+    @staticmethod
+    def _seeds_cooperative_exchange_3(minibatch):
+        nodes = {
+            ntype: [typed_seeds]
+            for ntype, typed_seeds in minibatch._seed_nodes.items()
+        }
+        minibatch._unique_future = unique_and_compact(
+            nodes, 0, 1, async_op=True
+        )
+        return minibatch
+
+    @staticmethod
+    def _seeds_cooperative_exchange_4(minibatch):
+        unique_seeds, inverse_seeds, _ = minibatch._unique_future.wait()
+        delattr(minibatch, "_unique_future")
+        inverse_seeds = {
+            ntype: typed_inv[0] for ntype, typed_inv in inverse_seeds.items()
+        }
+        minibatch._seed_nodes = revert_to_homo(unique_seeds)
+        sizes = {
+            ntype: typed_seeds.size(0)
+            for ntype, typed_seeds in unique_seeds.items()
+        }
+        minibatch._seed_sizes = revert_to_homo(sizes)
+        minibatch._seed_inverse_ids = revert_to_homo(inverse_seeds)
+        return minibatch
+
+    def _sample(self, minibatch):
+        (
+            minibatch.input_nodes,
+            minibatch.sampled_subgraphs,
+        ) = self.sample_subgraphs(
+            minibatch._seed_nodes, minibatch._seeds_timestamp
+        )
+        return minibatch
+
+    def sampling_stages(self, datapipe):
+        """The sampling stages are defined here by chaining to the datapipe. The
+        default implementation expects :meth:`sample_subgraphs` to be
+        implemented. To define fine-grained stages, this method should be
+        overridden.
+        """
+        return datapipe.transform(self._sample)
+
+    @staticmethod
+    def _seeds_preprocess(
+        minibatch: MiniBatch,
+        rank: int = 0,
+        world_size: int = 1,
+        async_op: bool = False,
+    ):
+        """Preprocess `seeds` in a minibatch to construct `unique_seeds`,
+        `node_timestamp` and `compacted_seeds` for further sampling. It
+        optionally incorporates timestamps for temporal graphs, organizing and
+        compacting seeds based on their types and timestamps. In heterogeneous
+        graph, `seeds` with same node type will be unqiued together.
+
+        Parameters
+        ----------
+        minibatch: MiniBatch
+            The minibatch.
+        rank : int
+            The rank of the current process among cooperating processes.
+        world_size : int
+            The number of cooperating
+            (`arXiv:2210.13339<https://arxiv.org/abs/2310.12403>`__) processes.
+        async_op: bool
+            Boolean indicating whether the call is asynchronous. If so, the
+            result can be obtained by calling wait on the returned future.
+
+        Returns
+        -------
+        unique_seeds: torch.Tensor or Dict[str, torch.Tensor]
+            A tensor or a dictionary of tensors representing the unique seeds.
+            In heterogeneous graphs, seeds are returned for each node type.
+        nodes_timestamp: None or a torch.Tensor or Dict[str, torch.Tensor]
+            Containing timestamps for each seed. This is only returned if
+            `minibatch` includes timestamps and the graph is temporal.
+        compacted_seeds: torch.tensor or a Dict[str, torch.Tensor]
+            Representation of compacted seeds corresponding to 'seeds', where
+            all node ids inside are compacted.
+        offsets: None or torch.Tensor or Dict[src, torch.Tensor]
+            The unique nodes offsets tensor partitions the unique_nodes tensor.
+            Has size `world_size + 1` and
+            `unique_nodes[offsets[i]: offsets[i + 1]]` belongs to the rank
+            `(rank + i) % world_size`.
+        """
+        use_timestamp = hasattr(minibatch, "timestamp")
+        assert (
+            not use_timestamp or world_size == 1
+        ), "Temporal code path does not currently support Cooperative Minibatching"
+        seeds = minibatch.seeds
+        is_heterogeneous = isinstance(seeds, Dict)
+        if is_heterogeneous:
+            # Collect nodes from all types of input.
+            nodes = defaultdict(list)
+            nodes_timestamp = None
+            if use_timestamp:
+                nodes_timestamp = defaultdict(list)
+            for seed_type, typed_seeds in seeds.items():
+                # When typed_seeds is a one-dimensional tensor, it represents
+                # seed nodes, which does not need to do unique and compact.
+                if typed_seeds.ndim == 1:
+                    nodes_timestamp = (
+                        minibatch.timestamp
+                        if hasattr(minibatch, "timestamp")
+                        else None
+                    )
+                    result = _NoOpWaiter((seeds, nodes_timestamp, None, None))
+                    break
+                result = None
+                assert typed_seeds.ndim == 2, (
+                    "Only tensor with shape 1*N and N*M is "
+                    + f"supported now, but got {typed_seeds.shape}."
+                )
+                ntypes = seed_type_str_to_ntypes(
+                    seed_type, typed_seeds.shape[1]
+                )
+                if use_timestamp:
+                    negative_ratio = (
+                        typed_seeds.shape[0]
+                        // minibatch.timestamp[seed_type].shape[0]
+                        - 1
+                    )
+                    neg_timestamp = minibatch.timestamp[
+                        seed_type
+                    ].repeat_interleave(negative_ratio)
+                for i, ntype in enumerate(ntypes):
+                    nodes[ntype].append(typed_seeds[:, i])
+                    if use_timestamp:
+                        nodes_timestamp[ntype].append(
+                            minibatch.timestamp[seed_type]
+                        )
+                        nodes_timestamp[ntype].append(neg_timestamp)
+
+            class _Waiter:
+                def __init__(self, nodes, nodes_timestamp, seeds):
+                    # Unique and compact the collected nodes.
+                    if use_timestamp:
+                        self.future = compact_temporal_nodes(
+                            nodes, nodes_timestamp
+                        )
+                    else:
+                        self.future = unique_and_compact(
+                            nodes, rank, world_size, async_op
+                        )
+                    self.seeds = seeds
+
+                def wait(self):
+                    """Returns the stored value when invoked."""
+                    if use_timestamp:
+                        unique_seeds, nodes_timestamp, compacted = self.future
+                        offsets = None
+                    else:
+                        unique_seeds, compacted, offsets = (
+                            self.future.wait() if async_op else self.future
+                        )
+                        nodes_timestamp = None
+                    seeds = self.seeds
+                    # Ensure there is no memory leak.
+                    self.future = self.seeds = None
+
+                    compacted_seeds = {}
+                    # Map back in same order as collect.
+                    for seed_type, typed_seeds in seeds.items():
+                        ntypes = seed_type_str_to_ntypes(
+                            seed_type, typed_seeds.shape[1]
+                        )
+                        compacted_seed = []
+                        for ntype in ntypes:
+                            compacted_seed.append(compacted[ntype].pop(0))
+                        compacted_seeds[seed_type] = (
+                            torch.cat(compacted_seed).view(len(ntypes), -1).T
+                        )
+
+                    return (
+                        unique_seeds,
+                        nodes_timestamp,
+                        compacted_seeds,
+                        offsets,
+                    )
+
+            # When typed_seeds is not a one-dimensional tensor
+            if result is None:
+                result = _Waiter(nodes, nodes_timestamp, seeds)
+        else:
+            # When seeds is a one-dimensional tensor, it represents seed nodes,
+            # which does not need to do unique and compact.
+            if seeds.ndim == 1:
+                nodes_timestamp = (
+                    minibatch.timestamp
+                    if hasattr(minibatch, "timestamp")
+                    else None
+                )
+                result = _NoOpWaiter((seeds, nodes_timestamp, None, None))
+            else:
+                # Collect nodes from all types of input.
+                nodes = [seeds.view(-1)]
+                nodes_timestamp = None
+                if use_timestamp:
+                    # Timestamp for source and destination nodes are the same.
+                    negative_ratio = (
+                        seeds.shape[0] // minibatch.timestamp.shape[0] - 1
+                    )
+                    neg_timestamp = minibatch.timestamp.repeat_interleave(
+                        negative_ratio
+                    )
+                    seeds_timestamp = torch.cat(
+                        (minibatch.timestamp, neg_timestamp)
+                    )
+                    nodes_timestamp = [
+                        seeds_timestamp for _ in range(seeds.shape[1])
+                    ]
+
+                class _Waiter:
+                    def __init__(self, nodes, nodes_timestamp, seeds):
+                        # Unique and compact the collected nodes.
+                        if use_timestamp:
+                            self.future = compact_temporal_nodes(
+                                nodes, nodes_timestamp
+                            )
+                        else:
+                            self.future = unique_and_compact(
+                                nodes, async_op=async_op
+                            )
+                        self.seeds = seeds
+
+                    def wait(self):
+                        """Returns the stored value when invoked."""
+                        if use_timestamp:
+                            (
+                                unique_seeds,
+                                nodes_timestamp,
+                                compacted,
+                            ) = self.future
+                            offsets = None
+                        else:
+                            unique_seeds, compacted, offsets = (
+                                self.future.wait() if async_op else self.future
+                            )
+                            nodes_timestamp = None
+                        seeds = self.seeds
+                        # Ensure there is no memory leak.
+                        self.future = self.seeds = None
+
+                        # Map back in same order as collect.
+                        compacted_seeds = compacted[0].view(seeds.shape)
+
+                        return (
+                            unique_seeds,
+                            nodes_timestamp,
+                            compacted_seeds,
+                            offsets,
+                        )
+
+                result = _Waiter(nodes, nodes_timestamp, seeds)
+
+        return result if async_op else result.wait()
+
+    def sample_subgraphs(
+        self, seeds, seeds_timestamp, seeds_pre_time_window=None
+    ):
+        """Sample subgraphs from the given seeds, possibly with temporal constraints.
+
+        Any subclass of SubgraphSampler should implement this method.
+
+        Parameters
+        ----------
+        seeds : Union[torch.Tensor, Dict[str, torch.Tensor]]
+            The seed nodes.
+
+        seeds_timestamp : Union[torch.Tensor, Dict[str, torch.Tensor]]
+            The timestamps of the seed nodes. If given, the sampled subgraphs
+            should not contain any nodes or edges that are newer than the
+            timestamps of the seed nodes. Default: None.
+
+        seeds_pre_time_window : Union[torch.Tensor, Dict[str, torch.Tensor]]
+            The time window of the nodes represents a period of time before
+            `seeds_timestamp`. If provided, only neighbors and related edges
+            whose timestamps fall within `[seeds_timestamp -
+            seeds_pre_time_window, seeds_timestamp]` will be filtered.
+        Returns
+        -------
+        Union[torch.Tensor, Dict[str, torch.Tensor]]
+            The input nodes.
+        List[SampledSubgraph]
+            The sampled subgraphs.
+
+        Examples
+        --------
+        >>> @functional_datapipe("my_sample_subgraph")
+        >>> class MySubgraphSampler(SubgraphSampler):
+        >>>     def __init__(self, datapipe, graph, fanouts):
+        >>>         super().__init__(datapipe)
+        >>>         self.graph = graph
+        >>>         self.fanouts = fanouts
+        >>>     def sample_subgraphs(self, seeds):
+        >>>         # Sample subgraphs from the given seeds.
+        >>>         subgraphs = []
+        >>>         subgraphs_nodes = []
+        >>>         for fanout in reversed(self.fanouts):
+        >>>             subgraph = self.graph.sample_neighbors(seeds, fanout)
+        >>>             subgraphs.insert(0, subgraph)
+        >>>             subgraphs_nodes.append(subgraph.nodes)
+        >>>             seeds = subgraph.nodes
+        >>>         subgraphs_nodes = torch.unique(torch.cat(subgraphs_nodes))
+        >>>         return subgraphs_nodes, subgraphs
+        """
+        raise NotImplementedError