Add typed-channel PPR C++ extraction support

gigl-core/core/sampling/ppr_forward_push.h

@@ -11,6 +11,16 @@
 
 namespace gigl {
 
+// Neighbor fetch input from Python, keyed by integer edge type ID:
+//   node_ids[N], flat_neighbor_ids[sum(counts)], counts[N].
+using NeighborFetchTensors = std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>;
+using NeighborFetchMap = std::unordered_map<int32_t, NeighborFetchTensors>;
+
+// PPR extraction output, keyed by integer node type ID:
+//   ids, weights/edge_attr, valid_counts.
+using PPRExtractTensors = std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>;
+using PPRExtractResult = std::unordered_map<int32_t, PPRExtractTensors>;
+
 // Per-seed, per-node-type PPR algorithm state.
 // Grouping all four tables into one struct is a logical convenience: a single
 // _state[seedIdx][nodeTypeId] access reaches all four tables for a given (seed, ntype)
@@ -25,6 +35,39 @@ struct SeedNodeTypeState {
     std::unordered_set<int32_t> queuedNodes;       // snapshot captured by drainQueue()
 };
 
+// Batched drain result for typed-PPR channels.
+//
+// Typed PPR keeps one PPRForwardPush state per channel.  During an iteration,
+// each channel drains its own queue, but Python should issue at most one shared
+// neighbor fetch per edge type. This struct carries both pieces of information:
+// 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. Channel IDs
+    // are positional indices into the states/channel-target vectors. Python
+    // builds those vectors from typed-channel insertion order, and this function
+    // appends indices in ascending order, so the ordering is stable.
+    //
+    // These channels need pushResiduals(), even when no fetch budget remains.
+    // In that case Python passes an empty fetch map and the channel uses its
+    // existing neighbor cache / budget-exhausted behavior, matching untyped PPR.
+    std::vector<int32_t> drainedChannelIndices;
+
+    // Subset of drainedChannelIndices that still have fetch budget and at least
+    // one non-empty uncached frontier.
+    std::vector<int32_t> fetchChannelIndices;
+
+    // Edge types requested by each fetch channel, aligned with fetchChannelIndices.
+    std::vector<std::vector<int32_t>> edgeTypeIdsByFetchChannel;
+
+    // Unioned node frontier for one shared distributed neighbor fetch. This is
+    // keyed by integer edge type ID, not node type ID, because neighbor fetches
+    // are edge-type scoped; node type alone would lose the relation/destination
+    // distinction for heterogeneous graphs. Tensor values are int64 source node
+    // IDs to fetch.
+    std::unordered_map<int32_t, torch::Tensor> unionedNodeIdsByEdgeTypeId;
+};
+
 // C++ kernel for PPR Forward Push (Andersen et al., 2006).
 // Hot-loop state lives here; distributed neighbor fetches are driven from Python.
 //
@@ -51,11 +94,7 @@ class PPRForwardPush {
     std::optional<std::unordered_map<int32_t, torch::Tensor>> drainQueue();
 
     // Push residuals given fetched neighbor data.
-    // fetchedByEtypeId: {etype_id: (node_ids[N], flat_nbrs[sum(counts)], counts[N])}
-    // TODO: Move these repeated tensor tuple/map types into aliases in a follow-up
-    // refactor-only PR.
-    void pushResiduals(
-        const std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>& fetchedByEtypeId);
+    void pushResiduals(const NeighborFetchMap& fetchedByEtypeId);
 
     // Return top-k PPR nodes plus residual-mass top-up nodes, sorted by score.
     //
@@ -72,8 +111,12 @@ class PPRForwardPush {
     // it is not a global top-k over ppr_score + residual when maxPPRNodes is tight.
     // maxPPRNodes is the final per-seed cap across finalized PPR and residual
     // top-up candidates.
-    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> extractTopKWithResidualTopUp(
-        int32_t maxPPRNodes, bool enableResidualTopUp);
+    PPRExtractResult extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp);
+
+    friend PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector<PPRForwardPush*>& states,
+                                                              const std::vector<int32_t>& channelTargetCounts,
+                                                              int32_t maxPPRNodes,
+                                                              bool enableResidualTopUp);
 
 private:
     // Total out-degree of a node across all edge types. Returns 0 for sink nodes.
@@ -116,4 +159,59 @@ class PPRForwardPush {
     std::unordered_map<uint64_t, std::vector<int32_t>> _neighborCache;
 };
 
+// Helper function for draining several independent channel states for one
+// typed-PPR iteration.
+//
+// This is the typed wrapper around PPRForwardPush::drainQueue(): it drains the
+// independent channel states, records every channel that still needs
+// pushResiduals(), and unions fetchable frontier nodes by edge type so Python
+// can issue one shared distributed neighbor fetch for duplicate channel
+// requests. Channel drains may run concurrently; the result merge happens in
+// channel order to keep the returned channel-index lists deterministic.
+//
+// Inputs:
+//   states: One PPRForwardPush per typed channel. Each state is mutated by its
+//           drainQueue() call.
+//   fetchIterationCounts: Number of distributed fetches already issued for each
+//                         channel; aligned with states.
+//   maxFetchIterations: -1 means unbounded; otherwise channels at this count
+//                       still need pushResiduals() but contribute no new fetch
+//                       frontier.
+//
+// Expected output: TypedPPRQueueDrainResult, whose drainedChannelIndices are
+// the channels to push this iteration and whose unionedNodeIdsByEdgeTypeId is
+// the shared fetch request.
+TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector<PPRForwardPush*>& states,
+                                                    const std::vector<int32_t>& fetchIterationCounts,
+                                                    int32_t maxFetchIterations);
+
+// Helper function for extracting and merging completed typed-PPR channel states
+// in one C++ step.
+//
+// For each seed/node-type, typed extraction builds one candidate view per
+// channel. When residual top-up is enabled, residual candidates are included in
+// that same view, so finalized PPR and residual top-up both obey the configured
+// channel target counts. The merge calibrates scores within each channel,
+// deduplicates candidates seen through multiple channels by attributing each
+// node to its highest-scoring channel, fills each channel target, redistributes
+// unused slots globally by score, and emits per-node-type tensors.
+//
+// Inputs:
+//   states: Completed PPRForwardPush states, one per typed channel.
+//   channelTargetCounts: Per-channel target output counts, aligned with states.
+//   maxPPRNodes: Maximum number of deduplicated nodes to return per seed.
+//   enableResidualTopUp: Whether residual candidates may participate in target
+//                        filling alongside finalized PPR candidates.
+//
+// Expected output: per-node-type tensors. Tuple values match
+// extractTopKWithResidualTopUp:
+//   ids: int64 node IDs, flattened across seeds.
+//   weights: double feature matrix with columns
+//            [best_calibrated_score, per-channel scores..., presence bits...].
+//   valid_counts: int64 count of selected nodes per seed.
+PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector<PPRForwardPush*>& states,
+                                                   const std::vector<int32_t>& channelTargetCounts,
+                                                   int32_t maxPPRNodes,
+                                                   bool enableResidualTopUp);
+
 } // namespace gigl

gigl-core/core/sampling/python_ppr_forward_push.cpp

@@ -21,7 +21,7 @@ namespace gigl {
 // pushResiduals receives Python-owned containers, so convert them while the GIL
 // is held and release only around the C++ state update.
 static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedByEtypeId) {
-    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> neighborTensorsByEtypeId;
+    NeighborFetchMap neighborTensorsByEtypeId;
     // Dict iteration touches Python objects — GIL must be held here.
     for (auto item : fetchedByEtypeId) {
         auto edgeTypeId = item.first.cast<int32_t>();
@@ -43,20 +43,21 @@ static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedB
 }
 
 static std::optional<std::unordered_map<int32_t, torch::Tensor>> drainQueueWrapper(PPRForwardPush& state) {
-    std::optional<std::unordered_map<int32_t, torch::Tensor>> drained;
+    std::optional<std::unordered_map<int32_t, torch::Tensor>> queueDrainResult;
     // drainQueue mutates only this PPRForwardPush instance and materializes CPU
     // tensors for frontier node IDs. pybind converts those tensor handles back
     // to Python tensors after return without copying the underlying storage.
     {
         py::gil_scoped_release release;
-        drained = state.drainQueue();
+        queueDrainResult = state.drainQueue();
     }
-    return drained;
+    return queueDrainResult;
 }
 
-static std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
-extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) {
-    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
+static PPRExtractResult extractTopKWithResidualTopUpWrapper(PPRForwardPush& state,
+                                                            int32_t maxPPRNodes,
+                                                            bool enableResidualTopUp) {
+    PPRExtractResult result;
     // Extraction walks C++ state and builds torch tensors. Returning through
     // pybind creates Python container/wrapper objects, not tensor data copies.
     {
@@ -66,6 +67,58 @@ extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes,
     return result;
 }
 
+static py::tuple drainTypedPPRChannelQueuesWrapper(const py::sequence& states,
+                                                   const std::vector<int32_t>& fetchIterationCounts,
+                                                   int32_t maxFetchIterations) {
+    std::vector<PPRForwardPush*> statePtrs;
+    statePtrs.reserve(py::len(states));
+    // Sequence iteration and casting touch Python objects, so keep the GIL
+    // while copying raw C++ state pointers out of the Python container.
+    for (py::handle stateObj : states) {
+        statePtrs.push_back(&stateObj.cast<PPRForwardPush&>());
+    }
+
+    // C++ typed drain only reads/mutates PPRForwardPush states and builds C++
+    // containers. Reacquire the GIL before constructing the Python tuple.
+    // REQUIREMENT: no other thread may read or mutate these channel states
+    // while the GIL is released. The typed sampler drains and pushes each
+    // channel in a single sequenced loop iteration.
+    TypedPPRQueueDrainResult queueDrainResult;
+    {
+        py::gil_scoped_release release;
+        queueDrainResult = drainTypedPPRChannelQueues(statePtrs, fetchIterationCounts, maxFetchIterations);
+    }
+    // Pybind converts the temporary C++ containers into Python objects. Tensor
+    // values are handles, so this does not copy tensor storage across the
+    // Python/C++ boundary.
+    return py::make_tuple(queueDrainResult.drainedChannelIndices,
+                          queueDrainResult.fetchChannelIndices,
+                          queueDrainResult.edgeTypeIdsByFetchChannel,
+                          queueDrainResult.unionedNodeIdsByEdgeTypeId);
+}
+
+static PPRExtractResult extractTypedTopKWithResidualTopUpWrapper(const py::sequence& states,
+                                                                 const std::vector<int32_t>& channelTargetCounts,
+                                                                 int32_t maxPPRNodes,
+                                                                 bool enableResidualTopUp) {
+    std::vector<PPRForwardPush*> statePtrs;
+    statePtrs.reserve(py::len(states));
+    // Sequence iteration and casting touch Python objects, so keep the GIL
+    // while copying raw C++ state pointers out of the Python container.
+    for (py::handle stateObj : states) {
+        statePtrs.push_back(&stateObj.cast<PPRForwardPush&>());
+    }
+
+    // C++ extraction only reads the completed channel states and builds C++
+    // tensors/containers. Reacquire the GIL before pybind converts the return.
+    PPRExtractResult result;
+    {
+        py::gil_scoped_release release;
+        result = extractTypedTopKWithResidualTopUp(statePtrs, channelTargetCounts, maxPPRNodes, enableResidualTopUp);
+    }
+    return result;
+}
+
 } // namespace gigl
 
 // TORCH_EXTENSION_NAME is set by PyTorch's build system to match the Python
@@ -85,7 +138,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         .def("drain_queue", gigl::drainQueueWrapper)
         .def("push_residuals", gigl::pushResidualsWrapper)
         .def("extract_top_k_with_residual_top_up",
-             gigl::extractTopKWithResidualTopUpWrapper,
+             &gigl::extractTopKWithResidualTopUpWrapper,
              py::arg("max_ppr_nodes"),
              py::arg("enable_residual_topup"));
+    m.def("drain_typed_ppr_channel_queues",
+          &gigl::drainTypedPPRChannelQueuesWrapper,
+          py::arg("states"),
+          py::arg("fetch_iteration_counts"),
+          py::arg("max_fetch_iterations") = -1);
+    m.def("extract_typed_top_k_with_residual_top_up",
+          &gigl::extractTypedTopKWithResidualTopUpWrapper,
+          py::arg("states"),
+          py::arg("channel_target_counts"),
+          py::arg("max_ppr_nodes"),
+          py::arg("enable_residual_topup"));
 }

gigl-core/src/gigl_core/__init__.py

@@ -1,3 +1,11 @@
-from gigl_core.ppr_forward_push import PPRForwardPush
+from gigl_core.ppr_forward_push import (
+    PPRForwardPush,
+    drain_typed_ppr_channel_queues,
+    extract_typed_top_k_with_residual_top_up,
+)
 
-__all__ = ["PPRForwardPush"]
+__all__ = [
+    "PPRForwardPush",
+    "drain_typed_ppr_channel_queues",
+    "extract_typed_top_k_with_residual_top_up",
+]

gigl-core/src/gigl_core/ppr_forward_push.pyi

@@ -1,3 +1,5 @@
+from typing import Sequence
+
 import torch
 
 class PPRForwardPush:
@@ -21,3 +23,15 @@ class PPRForwardPush:
         max_ppr_nodes: int,
         enable_residual_topup: bool,
     ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ...
+
+def drain_typed_ppr_channel_queues(
+    states: Sequence[PPRForwardPush],
+    fetch_iteration_counts: Sequence[int],
+    max_fetch_iterations: int = -1,
+) -> tuple[list[int], list[int], list[list[int]], dict[int, torch.Tensor]]: ...
+def extract_typed_top_k_with_residual_top_up(
+    states: Sequence[PPRForwardPush],
+    channel_target_counts: Sequence[int],
+    max_ppr_nodes: int,
+    enable_residual_topup: bool,
+) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ...