Add typed-channel PPR sampler by mkolodner-sc · Pull Request #709 · Snapchat/GiGL

mkolodner-sc · 2026-07-15T18:55:01Z

Summary

Adds the Python/API layer for typed-channel PPR sampling.

This PR is stacked on #724, which contains the C++ typed PPR drain/extraction implementation. This layer exposes those C++ APIs through pybind/stubs, adds the sampler option, and wires typed PPR execution into DistPPRNeighborSampler.

Motivation

For heterogeneous graphs, edge types often carry meaningful relation information. A single collapsed PPR traversal can tell us that a neighbor is relevant, but it loses which typed relation paths contributed to that relevance.

Typed-channel PPR preserves that structure by running PPR over configured edge-type traversal channels and emitting channel-aware edge attributes. Downstream models can consume the best calibrated PPR score, per-channel calibrated scores, and per-channel presence indicators.

What Changed

Adds typed_channel_quotas to PPRSamplerOptions.
- Keys can be a single canonical edge type, e.g. ("user", "views", "item").
- Keys can also be a tuple of canonical edge types to define one grouped traversal channel.
- Values are positive candidate quotas for that channel.
Adds typed-channel parsing, validation, and traversal-map construction.
- Rejects invalid channel keys and non-positive quotas.
- Rejects edge types that are not traversable by PPR.
- Restricts typed PPR to heterogeneous sampling.
- Allows quotas to sum above max_ppr_nodes so sparse channels can still help fill the final sequence.
Wires per-channel PPR execution into the distributed sampler.
- Builds one PPRForwardPush state per typed channel.
- Each state traverses only the edge types allowed by that channel.
- Calls the C++ typed drain/extraction APIs from Add typed-channel PPR C++ extraction support #724.
- Runs typed drain, push, and extraction work off the asyncio event-loop thread.
Exposes typed PPR outputs through pybind and Python stubs.
- Untyped PPR keeps scalar edge attrs.
- Typed PPR uses multi-column edge attrs from Add typed-channel PPR C++ extraction support #724: [best_score, channel_scores..., channel_presence_bits...].

Tests

Added focused Python coverage for:

typed-channel quota parsing and validation
traversal-map construction for typed channels
disabling residual top-up for typed PPR
sampler wiring around typed-channel options

kmontemayor2-sc

Thanks Matt! Took a first pass here :)

kmontemayor2-sc · 2026-07-29T14:51:44Z

+            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


BTW any reason we chose the "typed channel" moniker vs "meta path" which I think I've seen to describe this "sequence of edge types" elsewhere?

A meta path implies an ordered sequence of edge types usually to my understanding. In this case, the channels don't necessarily have an ordering -- they are just edges that we are permitted to traverse across for that forward push operation. As a result I've named them differently here.

kmontemayor2-sc · 2026-07-29T14:53:04Z

+            Quotas may sum above ``max_ppr_nodes`` to give sparse or
+            overlapping channels room to fill the sequence.


Sorry what does this mean? Can we give examples?

Added examples, this is saying quotas are candidate budgets, not guaranteed output counts. A channel may find fewer candidates than its quota, and multiple channels may find the same node, so the merged/deduped result can be smaller than the sum of quotas. We still cap final output at max_ppr_nodes, but allowing quotas to sum above that cap gives each channel room to contribute enough candidates before the global merge.

kmontemayor2-sc

Hi Matt, thanks for breaking up the PRs.

Does this comment make any sense to you? https://github.com/Snapchat/GiGL/pull/709/changes#r3677995032

I wonder if we want to think about the semantics here a bit more

kmontemayor2-sc · 2026-07-29T20:39:11Z

+    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=())
+


Should this test be in this PR?

kmontemayor2-sc · 2026-07-29T20:47:35Z

+# 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 quota order.


What does "quota order" mean here?

kmontemayor2-sc · 2026-07-29T21:03:34Z

+    async def _compute_typed_ppr_scores(
+        self,
+        seed_nodes: torch.Tensor,
+        seed_node_type: NodeType,
+        typed_ppr_channel_quotas: list[int],
+    ) -> HeteroPPRResult:
+        """Run one PPR state per typed channel and extract the merged result.
+
+        Each channel receives the same seed nodes but a different edge-type
+        traversal allowlist. Fetch frontiers are unioned across active channels
+        per iteration so shared graph neighborhoods are fetched once and reused
+        by every channel that requested them. After convergence, C++ applies
+        channel quotas, residual top-up, and cross-channel deduplication in one
+        extraction step.
+
+        Args:
+            seed_nodes: Global node IDs for the seed batch.
+            seed_node_type: Heterogeneous node type for ``seed_nodes``.
+            typed_ppr_channel_quotas: Per-channel candidate quotas, aligned
+                with ``self._typed_ppr_channel_to_node_type_id_to_edge_type_ids``.
+
+        Returns:
+            Heterogeneous PPR extraction output with typed edge-attribute
+            feature vectors.
+        """


btw I was sort of hoping these changes could be in a separate, non-api focussed PR.

(e.g. we'd throw if the typed_channel gets provided when we have the API but not logic behind it).

Sure, I'll migrate this to the other code which focuses on the C++/sampling logic changes.

kmontemayor2-sc · 2026-07-29T21:16:13Z

+// which channel states still need pushResiduals(), and the unioned frontier to
+// fetch once for all channels that requested it.
+struct TypedPPRQueueDrainResult {
+    // Channels whose drainQueue() returned a value this iteration.


typed channel state vector, aligned with channelQuotas and the Python typed_channel_quotas insertion order

Hmmm, do we ever need to worry about cross process serialization here? We had some issues with mangled order right? Was that just with proto?

mkolodner-sc requested review from kmontemayor2-sc, nshah-sc, svij-sc, xgao4-sc, yliu2-sc and zfan3-sc as code owners July 15, 2026 18:55

mkolodner-sc mentioned this pull request Jul 15, 2026

Move existing PPR sampler work off event loop #708

Merged

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 9724f59 to 81812c7 Compare July 15, 2026 20:02

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from fc410ab to 2b85456 Compare July 15, 2026 20:02

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 81812c7 to 41a65fd Compare July 15, 2026 21:24

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from 2b85456 to ecb394c Compare July 15, 2026 21:24

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 41a65fd to b826832 Compare July 15, 2026 21:40

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch 2 times, most recently from 0047b2f to 016517b Compare July 15, 2026 22:49

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch 2 times, most recently from f6fc7a8 to 65d131c Compare July 15, 2026 23:01

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from 016517b to 60a0cea Compare July 15, 2026 23:01

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 65d131c to 44307b8 Compare July 15, 2026 23:04

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch 2 times, most recently from ff2d3ea to d0f9cea Compare July 16, 2026 18:46

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch 2 times, most recently from 4d70c04 to a3751bd Compare July 16, 2026 19:15

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch 2 times, most recently from 33c8dde to 1b0ea2a Compare July 16, 2026 19:21

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch 2 times, most recently from 6fbb8c9 to ce171bb Compare July 16, 2026 20:01

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from 1b0ea2a to ed7387c Compare July 16, 2026 20:01

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from ed7387c to 457c30b Compare July 27, 2026 10:20

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from ce171bb to 52fcfbd Compare July 27, 2026 11:27

Base automatically changed from mkolodner-sc/pr701-ppr-event-loop to main July 27, 2026 11:50

mkolodner-sc force-pushed the mkolodner-sc/ppr-selector-helper branch from 2c09892 to 7b305fc Compare July 28, 2026 10:46

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from d2467d9 to 2b9ff1d Compare July 28, 2026 10:49

mkolodner-sc force-pushed the mkolodner-sc/ppr-selector-helper branch from 7b305fc to b1487ff Compare July 28, 2026 10:51

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 2b9ff1d to a105107 Compare July 28, 2026 10:54

mkolodner-sc force-pushed the mkolodner-sc/ppr-selector-helper branch from b1487ff to b5b0bb0 Compare July 28, 2026 10:55

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from a105107 to 98421e0 Compare July 28, 2026 10:56

mkolodner-sc force-pushed the mkolodner-sc/ppr-selector-helper branch 2 times, most recently from 3e093f5 to 930569f Compare July 28, 2026 11:18

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch 5 times, most recently from f2dcecb to ebc7ac7 Compare July 29, 2026 10:21

kmontemayor2-sc reviewed Jul 29, 2026

View reviewed changes

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from b0b2581 to 3e3d741 Compare July 29, 2026 16:13

mkolodner-sc mentioned this pull request Jul 29, 2026

Add typed-channel PPR C++ extraction support #724

Draft

mkolodner-sc changed the base branch from mkolodner-sc/ppr-selector-helper to mkolodner-sc/typed-channel-ppr-cpp July 29, 2026 16:21

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-cpp branch from b8bca9e to 41b08d4 Compare July 29, 2026 16:48

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from 3e3d741 to b0ef2fe Compare July 29, 2026 16:50

kmontemayor2-sc reviewed Jul 29, 2026

View reviewed changes

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch 7 times, most recently from 8e4520e to b326797 Compare July 30, 2026 08:04

mkolodner-sc added 2 commits July 30, 2026 08:21

Expose typed PPR sampler API

dc14ea3

Keep PPR result aliases in sampler

eb2ffe6

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from b326797 to eb2ffe6 Compare July 30, 2026 08:26

		Quotas may sum above ``max_ppr_nodes`` to give sparse or
		overlapping channels room to fill the sequence.

Uh oh!

Conversation

mkolodner-sc commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

What Changed

Tests

Uh oh!

kmontemayor2-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kmontemayor2-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mkolodner-sc commented Jul 15, 2026 •

edited

Loading