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 and extraction implementation. This layer exposes typed-channel PPR through sampler options and wires it into DistPPRNeighborSampler.

Motivation

A collapsed heterogeneous PPR traversal can identify relevant neighbors, but it loses which relation paths contributed to that relevance. Typed-channel PPR keeps that information by running PPR over configured edge-type channels and emitting channel-aware edge attributes.

Downstream consumers can use:

the best calibrated PPR score
per-channel calibrated scores
per-channel presence indicators

What Changed

Keeps max_ppr_nodes as the per-seed PPR sequence cap.
- For untyped PPR, this preserves the existing behavior.
- For typed PPR, this remains the total per-seed cap across all channels.
Adds typed_channel_ratios to PPRSamplerOptions.
- If omitted, the sampler uses the regular untyped PPR path.
- 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 ratios in (0, 1] and must sum to 1.0.
Converts typed-channel ratios to integer channel target counts during sampler construction.
- For example, max_ppr_nodes=200 and ratios {views: 0.6, likes/shares: 0.4} targets 120 and 80 nodes.
- Fractional counts are rounded deterministically while preserving the total max_ppr_nodes cap.
Adds typed-channel parsing, validation, and traversal-map construction.
- Restricts typed PPR to heterogeneous sampling.
- Rejects invalid channel keys and invalid ratio values.
- Builds per-channel edge traversal allowlists for the C++ PPR states.
Wires typed PPR execution into the distributed sampler.
- Builds one PPRForwardPush state per typed channel.
- Each state traverses only the edge types configured for that channel.
- Calls the typed drain and extraction APIs from Add typed-channel PPR C++ extraction support #724.
- Keeps typed PPR work off the asyncio event-loop thread where applicable.
Emits typed PPR edge attributes with shape [N, 1 + 2 * num_channels]:
- column 0: best score across channels
- next num_channels columns: per-channel scores
- final num_channels columns: per-channel presence bits

Tests

Added focused Python coverage for:

typed-channel ratio parsing and validation
ratio-to-target-count conversion
traversal-map construction
typed sampler option wiring
disabling residual top-up for typed PPR
end-to-end typed sampler output shape

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?

I'd prefer to keep this here, since this requires E2E (including API) changes to work and this branches out of the CPP PR. I plan to squash this in to the CPP PR and merge them together anyway.

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?

Discussed offline, we've reframed quota to be a decimal between 0 and 1 per channel for how much of the final sequence length they should be responsible for populating

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 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 10 times, most recently from e239141 to eb2ffe6 Compare August 3, 2026 05:54

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-cpp branch from a91cce7 to 006f24e Compare August 3, 2026 05:54

mkolodner-sc added 3 commits August 3, 2026 06:01

Expose typed PPR sampler API

f9e35e6

Keep PPR result aliases in sampler

1dd2d44

Use sequence length for typed PPR targets

c04b7cf

mkolodner-sc force-pushed the mkolodner-sc/typed-channel-ppr-sampler branch from eb2ffe6 to c04b7cf Compare August 3, 2026 06:07

mkolodner-sc added 2 commits August 3, 2026 15:52

Keep PPR sampler API on max_ppr_nodes

ab9ec35

Restore typed PPR ratio API

6d2a1a7

		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!

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!

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