diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index bcffb5ff0..48c7c8554 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/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, diff --git a/gigl/distributed/sampler_options.py b/gigl/distributed/sampler_options.py index b093b8e10..82756526a 100644 --- a/gigl/distributed/sampler_options.py +++ b/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] diff --git a/gigl/distributed/utils/dist_sampler.py b/gigl/distributed/utils/dist_sampler.py index e0289b26c..f73ebd189 100644 --- a/gigl/distributed/utils/dist_sampler.py +++ b/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: diff --git a/gigl/distributed/utils/dist_typed_sampler.py b/gigl/distributed/utils/dist_typed_sampler.py new file mode 100644 index 000000000..fbc32fdf1 --- /dev/null +++ b/gigl/distributed/utils/dist_typed_sampler.py @@ -0,0 +1,215 @@ +"""Construction-time typed-PPR option parsing for distributed samplers. + +The helpers in this module validate typed-channel edge-type keys, convert +public edge-type keys into the compact integer traversal maps consumed by the +C++ forward-push kernel, and derive integer channel target counts from public +ratios. They run once during ``DistPPRNeighborSampler`` initialization, not in +the per-batch PPR sampling hot loop. +""" + +import math +from collections.abc import Sequence +from typing import Optional, Union, cast + +from graphlearn_torch.typing import EdgeType, NodeType + +# Public typed_channel_ratios keys can be a single edge type or a grouped +# channel containing multiple edge types. +TypedPPRChannelKey = Union[EdgeType, tuple[EdgeType, ...]] + +"""TypedPPRChannelKey describes one public typed-PPR traversal channel key. + +A single canonical edge type creates one channel restricted to that edge type. +A tuple of canonical edge types creates one channel whose forward-push state may +traverse any edge type in the group. When typed PPR emits multi-column +``edge_attr`` tensors, channel columns follow the insertion order of the +``typed_channel_ratios`` mapping. +""" +# Parsed typed-channel edge-type allowlists, ordered to match the insertion +# order of typed_channel_ratios. +TypedPPRChannelEdgeTypeGroups = list[tuple[EdgeType, ...]] +# One channel's traversal map. The outer list is indexed by integer node-type +# ID; each inner list contains the integer edge-type IDs that channel may +# traverse from that node type. +TypedPPRChannelTraversalMap = list[list[int]] +# All typed-channel traversal maps, ordered to match typed-channel order. +TypedPPRChannelTraversalMaps = list[TypedPPRChannelTraversalMap] + + +def parse_typed_channel_ratio_groups( + typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]], +) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[float]]]: + """Parse typed-PPR channel keys and split keys from ratios. + + Public options allow each channel key to be either one canonical edge type + or a non-empty tuple of canonical edge types. Internally, traversal setup + needs only the edge-type groups while merge selection needs the aligned + ratio values, so this helper returns those two parallel lists. + + This is construction-time option parsing and is not part of the per-batch + PPR sampling hot loop. + + Args: + typed_channel_ratios: User-provided channel mapping from edge-type + allowlist to target output ratio. + + Returns: + ``(None, None)`` when typed PPR is disabled. Otherwise returns + ``(typed_channel_groups, typed_channel_ratio_list)``, both ordered by + the input mapping insertion order. + + Raises: + ValueError: If a channel key is not a canonical edge type or non-empty + tuple of canonical edge types, or if ratios are not positive and + summing to 1.0. + """ + if not typed_channel_ratios: + return None, None + + typed_channel_groups: TypedPPRChannelEdgeTypeGroups = [] + typed_channel_ratio_list: list[float] = [] + + def is_canonical_edge_type(value: object) -> bool: + """Return whether ``value`` has PyG's canonical edge-type shape.""" + return ( + isinstance(value, tuple) + and len(value) == 3 + and all(isinstance(part, str) for part in value) + ) + + for edge_type_key, ratio in typed_channel_ratios.items(): + if is_canonical_edge_type(edge_type_key): + edge_types = (cast(EdgeType, edge_type_key),) + elif ( + isinstance(edge_type_key, tuple) + and edge_type_key + and all(is_canonical_edge_type(edge_type) for edge_type in edge_type_key) + ): + edge_types = cast(tuple[EdgeType, ...], edge_type_key) + else: + raise ValueError( + "typed_channel_ratios keys must be a canonical edge type " + "(src_type, relation, dst_type) or a non-empty tuple of " + f"canonical edge types, got {edge_type_key!r}." + ) + if isinstance(ratio, bool) or ratio <= 0.0 or ratio > 1.0: + raise ValueError( + "typed_channel_ratios values must be positive ratios in (0, 1], " + f"got {ratio!r} for channel {edge_type_key!r}." + ) + typed_channel_groups.append(edge_types) + typed_channel_ratio_list.append(float(ratio)) + + ratio_sum = sum(typed_channel_ratio_list) + if not math.isclose(ratio_sum, 1.0, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError( + "typed_channel_ratios values must sum to 1.0, " + f"got {ratio_sum} from ratios {typed_channel_ratio_list}." + ) + + return typed_channel_groups, typed_channel_ratio_list + + +def compute_typed_channel_target_counts( + typed_channel_ratios: list[float], + max_ppr_nodes: int, +) -> list[int]: + """Convert typed-channel ratios to integer per-channel target counts. + + Ratios describe the desired attribution mix in the returned PPR sequence. + This helper converts them to integer counts whose sum is ``max_ppr_nodes``. + Fractional remainders are assigned from largest to smallest, with channel + order as the deterministic tie-breaker. + + This is construction-time option parsing and is not part of the per-batch + PPR sampling hot loop. + + Args: + typed_channel_ratios: Per-channel ratios, ordered by typed-channel + insertion order and summing to 1.0. + max_ppr_nodes: Maximum PPR sequence length per seed. + + Returns: + Integer target counts aligned with ``typed_channel_ratios``. + """ + raw_target_counts = [ratio * max_ppr_nodes for ratio in typed_channel_ratios] + target_counts = [math.floor(raw_count) for raw_count in raw_target_counts] + remaining_count = max_ppr_nodes - sum(target_counts) + channels_by_fractional_remainder = sorted( + range(len(raw_target_counts)), + key=lambda channel_index: ( + raw_target_counts[channel_index] - target_counts[channel_index], + -channel_index, + ), + reverse=True, + ) + for channel_index in channels_by_fractional_remainder[:remaining_count]: + target_counts[channel_index] += 1 + return target_counts + + +def build_edge_type_channel_group_edge_type_ids( + edge_type_groups: TypedPPRChannelEdgeTypeGroups, + edge_type_to_edge_type_id: dict[EdgeType, int], + node_type_to_edge_types: dict[NodeType, list[EdgeType]], + node_types: Sequence[NodeType], +) -> TypedPPRChannelTraversalMaps: + """Convert typed-channel edge-type allowlists to PPRForwardPush IDs. + + Returns one traversal map per typed channel, ordered to match + ``edge_type_groups``. A single traversal map has shape + ``list[list[int]]``: the outer index is ``node_type_id``, and the inner list + contains allowed ``edge_type_id`` values for that node type in that channel. + + This conversion runs once during sampler construction; the resulting integer + maps are reused by the per-batch C++ PPR states. + + Args: + edge_type_groups: Ordered typed channels, where each channel is the + canonical edge types that its PPR state may traverse. + edge_type_to_edge_type_id: Mapping from canonical edge type to the + compact integer ID used by the C++ forward-push kernel. + node_type_to_edge_types: Traversable edge types keyed by anchor node + type, after label-edge filtering and edge-direction handling. + node_types: Ordered node types whose positions match the kernel's + integer node-type IDs. + + Returns: + ``channel_traversal_maps[channel_id][node_type_id]`` gives the allowed + integer edge-type IDs that channel may traverse from that node type. + + Raises: + ValueError: If a configured edge type is unknown, excluded from PPR + traversal, or cannot be traversed from any node type. + """ + known_edge_types = set(edge_type_to_edge_type_id.keys()) + channel_edge_type_ids_by_node_type: TypedPPRChannelTraversalMaps = [] + for channel_edge_types in edge_type_groups: + unknown_edge_types = set(channel_edge_types) - known_edge_types + if unknown_edge_types: + raise ValueError( + "typed_channel_ratios includes non-traversable edge types " + f"{sorted(unknown_edge_types)!r}. Known traversable edge " + f"types are {sorted(known_edge_types)!r}." + ) + + channel_edge_type_set = set(channel_edge_types) + node_type_id_to_channel_edge_type_ids: TypedPPRChannelTraversalMap = [] + for node_type in node_types: + channel_edge_type_ids_for_node_type: list[int] = [] + for edge_type in node_type_to_edge_types.get(node_type, []): + if edge_type in channel_edge_type_set: + channel_edge_type_ids_for_node_type.append( + edge_type_to_edge_type_id[edge_type] + ) + node_type_id_to_channel_edge_type_ids.append( + channel_edge_type_ids_for_node_type + ) + if not any(node_type_id_to_channel_edge_type_ids): + raise ValueError( + "typed_channel_ratios includes edge-type " + f"channel={channel_edge_types!r}, " + "but no traversable edge types exist for that channel." + ) + channel_edge_type_ids_by_node_type.append(node_type_id_to_channel_edge_type_ids) + return channel_edge_type_ids_by_node_type diff --git a/tests/unit/distributed/dist_ppr_sampler_test.py b/tests/unit/distributed/dist_ppr_sampler_test.py index 604e650c6..6b79d041f 100644 --- a/tests/unit/distributed/dist_ppr_sampler_test.py +++ b/tests/unit/distributed/dist_ppr_sampler_test.py @@ -302,6 +302,47 @@ def _assert_ppr_scores_match_reference( ) +def _assert_typed_ppr_edge_attrs( + datum: HeteroData, + seed_type: str, + node_types: list[str], + num_channels: int, +) -> None: + """Assert typed PPR edge attributes have the expected channel layout.""" + expected_width = 1 + (2 * num_channels) + for node_type in node_types: + ppr_edge_type = (seed_type, "ppr", node_type) + assert ppr_edge_type in datum.edge_types, ( + f"Missing PPR edge type {ppr_edge_type} on HeteroData" + ) + ppr_edge_index = datum[ppr_edge_type].edge_index + ppr_edge_attr = datum[ppr_edge_type].edge_attr + + assert ppr_edge_index.dim() == 2 and ppr_edge_index.size(0) == 2, ( + f"Expected [2, X] edge_index, got shape {list(ppr_edge_index.shape)}" + ) + assert ppr_edge_attr.dim() == 2, ( + f"Expected 2D typed PPR edge_attr, got {ppr_edge_attr.dim()}D" + ) + assert ppr_edge_attr.size(0) == ppr_edge_index.size(1) + assert ppr_edge_attr.size(1) == expected_width, ( + f"Expected typed PPR edge_attr width {expected_width}, " + f"got {ppr_edge_attr.size(1)}" + ) + + if ppr_edge_attr.size(0) == 0: + continue + + best_scores = ppr_edge_attr[:, 0] + assert (best_scores >= 0).all() + assert (best_scores <= 1).all() + if best_scores.size(0) > 1: + assert (best_scores[:-1] >= best_scores[1:]).all() + + channel_presence = ppr_edge_attr[:, 1 + num_channels :] + assert ((channel_presence == 0) | (channel_presence == 1)).all() + + # --------------------------------------------------------------------------- # Spawned process functions # --------------------------------------------------------------------------- @@ -487,6 +528,82 @@ def _run_ppr_hetero_loader_correctness_check( shutdown_rpc() +def _run_typed_ppr_loader_shape_check(_: int) -> None: + """Verify typed PPR runs through DistNeighborLoader into HeteroData.""" + create_test_process_group() + + dataset = create_heterogeneous_dataset( + edge_indices=_TEST_HETERO_EDGE_INDICES, + edge_dir="out", + ) + node_ids = dataset.node_ids + assert isinstance(node_ids, dict) + + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(USER, node_ids[USER]), # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + num_neighbors=[], + sampler_options=PPRSamplerOptions( + alpha=_TEST_ALPHA, + eps=_TEST_EPS, + max_ppr_nodes=_TEST_MAX_PPR_NODES, + typed_channel_ratios={ + USER_TO_STORY: 0.5, + STORY_TO_USER: 0.5, + }, + ), + pin_memory_device=torch.device("cpu"), + batch_size=1, + ) + + batches_checked = 0 + for datum in loader: + assert isinstance(datum, HeteroData) + _assert_typed_ppr_edge_attrs( + datum=datum, + seed_type=str(USER), + node_types=[USER, STORY], + num_channels=2, + ) + for edge_type in datum.edge_types: + assert edge_type[1] == "ppr", ( + f"Non-PPR edge type {edge_type} found in PPR sampler output" + ) + batches_checked += 1 + assert batches_checked == _NUM_TEST_USERS, ( + f"Expected {_NUM_TEST_USERS} batches, got {batches_checked}" + ) + + empty_story_loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(USER, node_ids[USER][:1]), # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + num_neighbors=[], + sampler_options=PPRSamplerOptions( + alpha=_TEST_ALPHA, + eps=_TEST_EPS, + max_ppr_nodes=_TEST_MAX_PPR_NODES, + typed_channel_ratios={ + STORY_TO_USER: 0.5, + (STORY_TO_USER,): 0.5, + }, + ), + pin_memory_device=torch.device("cpu"), + batch_size=1, + ) + empty_story_datum = next(iter(empty_story_loader)) + assert isinstance(empty_story_datum, HeteroData) + _assert_typed_ppr_edge_attrs( + datum=empty_story_datum, + seed_type=str(USER), + node_types=[USER, STORY], + num_channels=2, + ) + empty_story_edge_attr = empty_story_datum[(str(USER), "ppr", STORY)].edge_attr + assert tuple(empty_story_edge_attr.shape) == (0, 5) + + shutdown_rpc() + + def _run_ppr_ablp_loader_correctness_check( _: int, alpha: float, @@ -814,6 +931,10 @@ def test_ppr_sampler_destination_only_node_type(self) -> None: """Verify PPR output includes destination-only node types.""" mp.spawn(fn=_run_ppr_destination_only_node_type, args=()) + def test_typed_ppr_sampler_loader_outputs_channel_attrs(self) -> None: + """Verify typed PPR runs end-to-end through the loader.""" + mp.spawn(fn=_run_typed_ppr_loader_shape_check, args=()) + def test_ppr_sampler_ablp_ignores_label_edges_for_anchor_ppr(self) -> None: """Verify ABLP label edges are excluded from anchor-seed PPR walks.""" mp.spawn(fn=_run_ppr_ablp_label_edges_do_not_affect_anchor_ppr, args=()) diff --git a/tests/unit/distributed/dist_typed_sampler_test.py b/tests/unit/distributed/dist_typed_sampler_test.py new file mode 100644 index 000000000..2cc765e92 --- /dev/null +++ b/tests/unit/distributed/dist_typed_sampler_test.py @@ -0,0 +1,84 @@ +"""Unit tests for typed-PPR sampler construction helpers.""" + +from absl.testing import absltest + +from gigl.distributed.utils.dist_typed_sampler import ( + build_edge_type_channel_group_edge_type_ids, + compute_typed_channel_target_counts, + parse_typed_channel_ratio_groups, +) +from tests.test_assets.distributed.test_dataset import ( + STORY, + STORY_TO_USER, + USER, + USER_TO_STORY, +) +from tests.test_assets.test_case import TestCase + + +class DistTypedSamplerTest(TestCase): + def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( + self, + ) -> None: + """Verify typed-PPR can use canonical edge-type channels.""" + node_type_to_edge_types = { + USER: [USER_TO_STORY], + STORY: [STORY_TO_USER], + } + node_types = [USER, STORY] + edge_type_to_edge_type_id = { + USER_TO_STORY: 0, + STORY_TO_USER: 1, + } + + typed_channel_groups, typed_channel_ratio_list = ( + parse_typed_channel_ratio_groups( + { + USER_TO_STORY: 0.6, + (USER_TO_STORY, STORY_TO_USER): 0.4, + } + ) + ) + assert typed_channel_groups is not None + assert typed_channel_ratio_list is not None + + self.assertEqual( + typed_channel_groups, + [ + (USER_TO_STORY,), + (USER_TO_STORY, STORY_TO_USER), + ], + ) + self.assertEqual(typed_channel_ratio_list, [0.6, 0.4]) + self.assertEqual( + compute_typed_channel_target_counts(typed_channel_ratio_list, 7), + [4, 3], + ) + self.assertEqual( + build_edge_type_channel_group_edge_type_ids( + edge_type_groups=typed_channel_groups, + edge_type_to_edge_type_id=edge_type_to_edge_type_id, + node_type_to_edge_types=node_type_to_edge_types, + node_types=node_types, + ), + [ + [[0], []], + [[0], [1]], + ], + ) + + with self.assertRaisesRegex(ValueError, "canonical edge type"): + parse_typed_channel_ratio_groups({("bad",): 1.0}) + with self.assertRaisesRegex(ValueError, "sum to 1.0"): + parse_typed_channel_ratio_groups({USER_TO_STORY: 0.5}) + with self.assertRaisesRegex(ValueError, "non-traversable edge types"): + build_edge_type_channel_group_edge_type_ids( + edge_type_groups=[(("unknown", "edge", "type"),)], + edge_type_to_edge_type_id=edge_type_to_edge_type_id, + node_type_to_edge_types=node_type_to_edge_types, + node_types=node_types, + ) + + +if __name__ == "__main__": + absltest.main()