Add typed-channel PPR sampler

gigl/distributed/dist_ppr_sampler.py

@@ -21,6 +21,13 @@
 from graphlearn_torch.utils import merge_dict
 
 from gigl.distributed.base_sampler import BaseDistNeighborSampler
+from gigl.distributed.utils.dist_typed_sampler import (
+    TypedPPRChannelKey,
+    TypedPPRChannelTraversalMaps,
+    build_edge_type_channel_group_edge_type_ids,
+    compute_typed_channel_target_counts,
+    parse_typed_channel_ratio_groups,
+)
 from gigl.types.graph import DEFAULT_HOMOGENEOUS_NODE_TYPE, is_label_edge_type
 
 # Trailing "." is an intentional separator.  These constants are used both to
@@ -40,6 +47,21 @@
     DEFAULT_HOMOGENEOUS_NODE_TYPE,
 )
 
+# C++ PPR extraction output: flat node IDs, flat weights, and per-seed valid
+# counts. Homogeneous extraction uses tensors directly; heterogeneous extraction
+# uses dictionaries keyed by node type.
+PPRResult = tuple[
+    Union[torch.Tensor, dict[NodeType, torch.Tensor]],
+    Union[torch.Tensor, dict[NodeType, torch.Tensor]],
+    Union[torch.Tensor, dict[NodeType, torch.Tensor]],
+]
+# Heterogeneous-only view of PPRResult after typed PPR extraction.
+HeteroPPRResult = tuple[
+    dict[NodeType, torch.Tensor],
+    dict[NodeType, torch.Tensor],
+    dict[NodeType, torch.Tensor],
+]
+
 
 class DistPPRNeighborSampler(BaseDistNeighborSampler):
     """Personalized PageRank (PPR) based distributed neighbor sampler.
@@ -87,7 +109,13 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler):
     **Heterogeneous (HeteroData)** — one PPR edge type per
     ``(seed_type, neighbor_type)`` pair, with ``"ppr"`` as the relation:
         - ``data[(seed_type, "ppr", neighbor_type)].edge_index``: same format as above.
-        - ``data[(seed_type, "ppr", neighbor_type)].edge_attr``: same format as above.
+        - ``data[(seed_type, "ppr", neighbor_type)].edge_attr``: scalar PPR
+          score for regular PPR. For typed PPR, edge attrs are multi-column:
+          ``[best_calibrated_score, calibrated_channel_scores..., channel_presence_bits...]``.
+          Typed-PPR scores are calibrated within each channel/seed pool and
+          globally ranked by the best calibrated score. Channel columns follow
+          the insertion order of ``typed_channel_ratios``. Column 0 is the
+          scalar best score for consumers that need a single PPR weight.
 
     Args:
         alpha: Restart probability (teleport probability back to seed). Higher values
@@ -104,6 +132,43 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler):
             during Forward Push when fewer than ``max_ppr_nodes`` finalized PPR
             scores are available.
         num_neighbors_per_hop: Maximum number of neighbors to fetch per hop.
+        typed_channel_ratios: Optional target proportions for typed PPR
+            traversal channels. If not provided, PPR uses the regular untyped path: each
+            state may traverse all eligible edge types for the current node
+            type and emits scalar PPR scores without channel attribution.
+            Keys may be either a single canonical edge type
+            ``(src_type, relation, dst_type)`` or a tuple of canonical edge
+            types. Each key defines one traversal channel whose PPR state may
+            traverse only those exact edge types.
+            Values are positive ratios that must sum to ``1.0``. The sampler
+            converts ratios to per-channel target counts from ``max_ppr_nodes``.
+            Finalized PPR candidates and residual top-up candidates both obey
+            these target counts. If the same node appears in multiple channels,
+            it is attributed to the channel where it has the highest calibrated
+            score. If sparse channels or duplicate nodes leave unused target
+            slots, the remaining slots are redistributed globally by score so
+            the returned sequence can still fill up to ``max_ppr_nodes``.
+            Example::
+
+                typed_channel_ratios = {
+                    ("user", "views", "item"): 0.6,
+                    (
+                        ("user", "likes", "item"),
+                        ("user", "shares", "item"),
+                    ): 0.4,
+                }
+
+            With ``max_ppr_nodes=200``, this example targets 120 nodes
+            attributed to the views channel and 80 nodes attributed to the
+            grouped likes/shares channel. The views channel traverses only
+            ``("user", "views", "item")`` edges. The grouped likes/shares
+            channel traverses either likes or shares edges in one PPR state.
+            Both finalized PPR rows and residual top-up rows fill those same
+            targets. The targets are best-effort rather than strict per-seed
+            guarantees: if a channel cannot provide enough unique candidates,
+            unused slots are filled by the remaining highest-scoring candidates
+            from any channel.
+
         degree_tensors: Pre-computed total-degree tensors (int32). Homogeneous
             graphs use a single tensor; heterogeneous graphs use tensors keyed
             by NodeType. The colocated and graph-store loader paths retrieve
@@ -121,6 +186,7 @@ def __init__(
         num_neighbors_per_hop: int = 100_000,
         degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]],
         max_fetch_iterations: Optional[int] = None,
+        typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]] = None,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
@@ -162,7 +228,22 @@ def __init__(
             ]
             self._is_homogeneous = True
 
-        self._typed_ppr_channel_target_counts: Optional[list[int]] = None
+        typed_channel_groups, typed_channel_ratio_list = (
+            parse_typed_channel_ratio_groups(typed_channel_ratios)
+        )
+        self._typed_ppr_channel_target_counts = (
+            compute_typed_channel_target_counts(
+                typed_channel_ratio_list,
+                max_ppr_nodes,
+            )
+            if typed_channel_ratio_list is not None
+            else None
+        )
+        if self._typed_ppr_channel_target_counts is not None:
+            if self._is_homogeneous:
+                raise ValueError(
+                    "Typed PPR channel ratios are only supported for heterogeneous PPR sampling."
+                )
 
         # Convert the public homogeneous/heterogeneous degree-tensor shape to
         # the node-type keyed form used internally by PPR.
@@ -226,9 +307,16 @@ def __init__(
             for node_type in all_node_types
         ]
 
-        self._typed_ppr_channel_to_node_type_id_to_edge_type_ids: list[
-            list[list[int]]
-        ] = []
+        self._typed_ppr_channel_to_node_type_id_to_edge_type_ids: TypedPPRChannelTraversalMaps = []
+        if typed_channel_groups is not None:
+            self._typed_ppr_channel_to_node_type_id_to_edge_type_ids = (
+                build_edge_type_channel_group_edge_type_ids(
+                    edge_type_groups=typed_channel_groups,
+                    edge_type_to_edge_type_id=self._etype_to_etype_id,
+                    node_type_to_edge_types=self._node_type_to_edge_types,
+                    node_types=self._ntype_id_to_ntype,
+                )
+            )
 
     def _convert_degree_tensors_to_dict(
         self,

gigl/distributed/sampler_options.py

@@ -13,6 +13,7 @@
 from graphlearn_torch.typing import EdgeType
 
 from gigl.common.logger import Logger
+from gigl.distributed.utils.dist_typed_sampler import TypedPPRChannelKey
 
 logger = Logger()
 
@@ -41,6 +42,10 @@ class PPRSamplerOptions:
     - ``edge_index``: ``[2, N]`` int64 — row 0 is local seed indices, row 1 is local
       neighbor indices.
     - ``edge_attr``: ``[N]`` float — PPR score for each (seed, neighbor) pair.
+      Typed PPR emits multi-column edge attrs:
+      ``[best_score, channel_scores..., channel_presence_bits...]``.
+      Column 0 is the scalar best score for consumers that need a single PPR
+      weight.
 
     For homogeneous graphs these live directly on ``data.edge_index`` / ``data.edge_attr``.
 
@@ -75,6 +80,49 @@ class PPRSamplerOptions:
             The algorithm still runs to convergence — re-enqueued nodes propagate
             through cached neighbors at negligible cost. ``None`` (default) means
             no fetch limit.
+        typed_channel_ratios: Optional target proportions for typed PPR
+            traversal channels defined by canonical edge-type allowlists. Keys
+            may be either a single canonical edge type
+            ``(src_type, relation, dst_type)`` or a tuple of canonical edge
+            types. Each key defines one traversal channel whose PPR state may
+            traverse only those exact edge types.
+            If not provided, PPR uses the regular untyped path: each state may
+            traverse all eligible edge types for the current node type and emits
+            scalar PPR scores without channel attribution.
+            Channel order follows the insertion order of this mapping. Values
+            are positive ratios that must sum to ``1.0``. The sampler converts
+            ratios to per-channel target counts from ``max_ppr_nodes``.
+            Finalized PPR candidates and residual top-up candidates both obey
+            these target counts. If the same node appears in multiple channels,
+            it is attributed to the channel where it has the highest calibrated
+            score. If sparse channels or duplicate nodes leave unused target
+            slots, the remaining slots are redistributed globally by score so
+            the returned sequence can still fill up to ``max_ppr_nodes``.
+            Example::
+
+                typed_channel_ratios = {
+                    ("user", "views", "item"): 0.6,
+                    (
+                        ("user", "likes", "item"),
+                        ("user", "shares", "item"),
+                    ): 0.4,
+                }
+
+            This example creates two traversal channels. The first channel can
+            traverse only ``("user", "views", "item")`` edges. With
+            ``max_ppr_nodes=200``, the ``0.6`` ratio targets 120 nodes
+            attributed to this channel. The second channel groups
+            ``("user", "likes", "item")`` and ``("user", "shares", "item")``
+            into one PPR state; the ``0.4`` ratio targets 80 nodes attributed
+            to that combined likes/shares channel. These targets are best-effort
+            rather than strict per-seed guarantees because channels may be
+            sparse or overlapping.
+
+            If residual top-up is enabled, discovered-but-unpushed residual
+            candidates from the same completed PPR states are included on the
+            same mass scale as finalized PPR scores: ``ppr_score + residual``.
+            Residual candidates follow the same channel targets as finalized
+            PPR candidates.
     """
 
     alpha: float = 0.5
@@ -83,6 +131,7 @@ class PPRSamplerOptions:
     enable_residual_topup: bool = True
     num_neighbors_per_hop: int = 1_000
     max_fetch_iterations: Optional[int] = None
+    typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]] = None
 
 
 SamplerOptions = Union[KHopNeighborSamplerOptions, PPRSamplerOptions]

gigl/distributed/utils/dist_sampler.py

@@ -85,6 +85,7 @@ def create_dist_sampler(
             enable_residual_topup=sampler_options.enable_residual_topup,
             max_fetch_iterations=sampler_options.max_fetch_iterations,
             num_neighbors_per_hop=sampler_options.num_neighbors_per_hop,
+            typed_channel_ratios=sampler_options.typed_channel_ratios,
             degree_tensors=degree_tensors,
         )
     else: