diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 915df7c56..5b5a16a76 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -147,8 +148,77 @@ std::optional> PPRForwardPush::drainQ return result; } -void PPRForwardPush::pushResiduals( - const std::unordered_map>& fetchedByEtypeId) { +TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector& states, + const std::vector& fetchIterationCounts, + int32_t maxFetchIterations) { + TORCH_CHECK(states.size() == fetchIterationCounts.size(), + "Expected one fetch iteration count per PPR state, got ", + fetchIterationCounts.size(), + " counts for ", + states.size(), + " states."); + + TypedPPRQueueDrainResult queueDrainResult; + // Fetch requests are edge-type scoped: each edge type has its own adjacency + // table and destination node type. Grouping by node type would merge + // relation-specific fetches that must remain separate. + std::unordered_map> unionedSourceNodeIdsByEdgeTypeId; + + std::vector>> frontiersByChannel; + frontiersByChannel.resize(states.size()); + if (states.size() == 1) { + frontiersByChannel[0] = states[0]->drainQueue(); + } else { + std::vector>>> frontierFutures; + frontierFutures.reserve(states.size()); + for (PPRForwardPush* state : states) { + frontierFutures.push_back(std::async(std::launch::async, [state]() { return state->drainQueue(); })); + } + for (size_t channelIndex = 0; channelIndex < frontierFutures.size(); ++channelIndex) { + frontiersByChannel[channelIndex] = frontierFutures[channelIndex].get(); + } + } + + // Merge per-channel drain results serially in channel order so the returned + // channel-index lists stay deterministic while the expensive independent + // drainQueue() calls above can run concurrently. + for (size_t channelIndex = 0; channelIndex < states.size(); ++channelIndex) { + auto& channelFrontierByEdgeTypeId = frontiersByChannel[channelIndex]; + if (!channelFrontierByEdgeTypeId.has_value()) { + continue; + } + + queueDrainResult.drainedChannelIndices.push_back(static_cast(channelIndex)); + + bool fetchBudgetRemaining = maxFetchIterations < 0 || fetchIterationCounts[channelIndex] < maxFetchIterations; + if (!fetchBudgetRemaining) { + continue; + } + + std::vector requestedEdgeTypeIds; + for (const auto& [edgeTypeId, nodes] : channelFrontierByEdgeTypeId.value()) { + requestedEdgeTypeIds.push_back(edgeTypeId); + auto nodeAccessor = nodes.accessor(); + auto& unionedSourceNodeIds = unionedSourceNodeIdsByEdgeTypeId[edgeTypeId]; + for (int64_t nodeIndex = 0; nodeIndex < nodes.size(0); ++nodeIndex) { + unionedSourceNodeIds.insert(nodeAccessor[nodeIndex]); + } + } + + if (!requestedEdgeTypeIds.empty()) { + queueDrainResult.fetchChannelIndices.push_back(static_cast(channelIndex)); + queueDrainResult.edgeTypeIdsByFetchChannel.push_back(std::move(requestedEdgeTypeIds)); + } + } + + for (const auto& [edgeTypeId, unionedSourceNodeIds] : unionedSourceNodeIdsByEdgeTypeId) { + std::vector nodeIdsToLookup(unionedSourceNodeIds.begin(), unionedSourceNodeIds.end()); + queueDrainResult.unionedNodeIdsByEdgeTypeId[edgeTypeId] = torch::tensor(nodeIdsToLookup, torch::kLong); + } + return queueDrainResult; +} + +void PPRForwardPush::pushResiduals(const NeighborFetchMap& fetchedByEtypeId) { // Step 1: Persist fetched neighbor lists in the per-state cache. drainQueue() // consults this cache before requesting future lookups, so storing every // fetched row here avoids re-fetching a (node, edge type) pair if it re-enters @@ -374,11 +444,221 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, } } -std::unordered_map> PPRForwardPush:: - extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) { +// Helper function for adding one channel's extracted PPR candidates into the +// emitted typed feature table and a selection candidate list. +// +// This helper writes the emitted edge_attr feature table for every node it +// sees. When residual top-up is enabled, callers pass residual-aware scores +// here so finalized and residual rows are emitted on the same ppr_score + +// residual scale. +// +// Inputs: +// nodesAndScores: extracted (node_id, score) candidates for one channel. +// maxScore: largest score in the channel, used to calibrate scores to [0, 1]. +// channelIndex: index of the typed PPR channel being added. +// numChannels: total typed PPR channels, used to derive feature width. +// outputScoresByNodeId: mutable map from node ID to emitted edge_attr features. +// channelOutputCandidates: mutable sortable candidates for target-count selection. +// +// Expected output: outputScoresByNodeId has this channel's score/presence +// features merged in, and channelOutputCandidates contains this channel's +// calibrated candidates. +static void addTypedPPRSeedFeaturesAndCandidates(const std::vector>& nodesAndScores, + double maxScore, + int32_t channelIndex, + int32_t numChannels, + std::unordered_map>& outputScoresByNodeId, + std::vector>& channelOutputCandidates) { + if (nodesAndScores.empty()) { + return; + } + TORCH_CHECK(maxScore > 0.0, + "Typed PPR output has candidates but non-positive max score ", + maxScore, + ", which indicates invalid PPR state."); + + // Feature width is 1 + 2C: + // column 0: best calibrated score across channels, used as the scalar + // PPR weight by downstream ranking/sequence construction. + // columns [1, C]: calibrated score for each channel. + // columns [1 + C, 1 + 2C): presence bit for each channel. + int32_t numEdgeAttrFeatures = 1 + (2 * numChannels); + for (const auto& [nodeId, score] : nodesAndScores) { + double calibratedScore = score / maxScore; + auto scoreIter = outputScoresByNodeId.find(nodeId); + if (scoreIter == outputScoresByNodeId.end()) { + scoreIter = outputScoresByNodeId.emplace(nodeId, std::vector(numEdgeAttrFeatures, 0.0)).first; + } + auto& scoreFeatures = scoreIter->second; + + // Feature layout: + // [best_calibrated_score, per-channel scores..., channel presence bits...]. + scoreFeatures[0] = std::max(scoreFeatures[0], calibratedScore); + int32_t channelScoreIndex = 1 + channelIndex; + int32_t channelPresenceIndex = 1 + numChannels + channelIndex; + + // Record this node's score for the current channel and mark that the + // channel reached the node. Current extraction emits one row per node + // per channel, so no intra-channel merge is needed here. + scoreFeatures[channelScoreIndex] = calibratedScore; + scoreFeatures[channelPresenceIndex] = 1.0; + + // Keep a per-channel sortable candidate list so target counts can be + // applied after cross-channel attribution. + channelOutputCandidates.emplace_back(nodeId, calibratedScore); + } +} + +// Helper function for applying typed channel target counts and cross-channel dedup. +// +// Inputs: +// candidatesByChannel: mutable (node_id, calibrated_score) candidates per channel. +// channelTargetCounts: desired output count per attributed channel. +// maxPPRNodes: maximum number of deduplicated node IDs to return for a seed. +// +// Expected output: node IDs selected by this policy: +// 1. Attribute duplicate nodes to the channel where they have the highest score. +// 2. Fill each channel up to its target count. +// 3. Redistribute unused slots to the remaining highest-scoring candidates. +// 4. Return the selected nodes globally ranked by best calibrated score. +static std::vector selectTypedPPRNodeIds( + std::vector>>& candidatesByChannel, + const std::vector& channelTargetCounts, + int32_t maxPPRNodes) { + struct AttributedCandidate { + double bestCalibratedScore; + double channelCalibratedScore; + int32_t channelIndex; + int32_t nodeId; + }; + + const auto higherAttributedCandidate = [](const auto& a, const auto& b) { + if (a.bestCalibratedScore != b.bestCalibratedScore) { + return a.bestCalibratedScore > b.bestCalibratedScore; + } + if (a.channelCalibratedScore != b.channelCalibratedScore) { + return a.channelCalibratedScore > b.channelCalibratedScore; + } + if (a.channelIndex != b.channelIndex) { + return a.channelIndex < b.channelIndex; + } + return a.nodeId < b.nodeId; + }; + + size_t totalCandidateRows = 0; + for (const auto& candidates : candidatesByChannel) { + totalCandidateRows += candidates.size(); + } + + // Deduplicate before applying channel targets. A node that appears in + // multiple channels is attributed to the channel where it has the strongest + // calibrated score, which is the same score used for global ranking. + std::unordered_map bestCandidateByNodeId; + bestCandidateByNodeId.reserve(totalCandidateRows); + for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByChannel.size()); ++channelIndex) { + for (const auto& [nodeId, calibratedScore] : candidatesByChannel[channelIndex]) { + AttributedCandidate candidate{calibratedScore, calibratedScore, channelIndex, nodeId}; + auto [candidateIter, inserted] = bestCandidateByNodeId.emplace(nodeId, candidate); + if (!inserted && higherAttributedCandidate(candidate, candidateIter->second)) { + candidateIter->second = candidate; + } + } + } + + std::vector> candidatesByAttributedChannel( + static_cast(candidatesByChannel.size())); + for (const auto& candidateEntry : bestCandidateByNodeId) { + const auto& candidate = candidateEntry.second; + candidatesByAttributedChannel[candidate.channelIndex].push_back(candidate); + } + + std::vector selectedCandidates; + selectedCandidates.reserve( + static_cast(std::min(maxPPRNodes, static_cast(bestCandidateByNodeId.size())))); + std::unordered_set selectedNodeIds; + selectedNodeIds.reserve(selectedCandidates.capacity()); + + // First honor the target counts for each attributed channel. Selection is + // local to each channel so we only spend target slots on that channel's + // strongest unique nodes. + for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByAttributedChannel.size()); + ++channelIndex) { + auto& candidates = candidatesByAttributedChannel[channelIndex]; + int32_t remainingOutputSlots = maxPPRNodes - static_cast(selectedCandidates.size()); + if (remainingOutputSlots <= 0) { + break; + } + int32_t numTargetCandidates = std::min( + {channelTargetCounts[channelIndex], static_cast(candidates.size()), remainingOutputSlots}); + if (numTargetCandidates <= 0) { + continue; + } + + if (numTargetCandidates < static_cast(candidates.size())) { + // Membership is enough here: the final selected set is sorted once + // after target fill and redistribution. + std::nth_element(candidates.begin(), + candidates.begin() + numTargetCandidates, + candidates.end(), + higherAttributedCandidate); + } + + for (int32_t candidateIndex = 0; candidateIndex < numTargetCandidates; ++candidateIndex) { + const auto& candidate = candidates[candidateIndex]; + selectedCandidates.push_back(candidate); + selectedNodeIds.insert(candidate.nodeId); + } + } + + // Sparse channels or cross-channel duplicates can leave some target slots + // unused. Redistribute that leftover capacity to the strongest remaining + // candidates globally so sequence length is not sacrificed for ratios that + // cannot be exactly filled on this seed. + int32_t remainingOutputSlots = maxPPRNodes - static_cast(selectedCandidates.size()); + if (remainingOutputSlots > 0) { + std::vector remainingCandidates; + remainingCandidates.reserve(bestCandidateByNodeId.size() - selectedNodeIds.size()); + for (const auto& candidateEntry : bestCandidateByNodeId) { + const auto& candidate = candidateEntry.second; + if (selectedNodeIds.find(candidate.nodeId) == selectedNodeIds.end()) { + remainingCandidates.push_back(candidate); + } + } + + int32_t numFillCandidates = std::min(remainingOutputSlots, static_cast(remainingCandidates.size())); + if (numFillCandidates < static_cast(remainingCandidates.size())) { + // Membership is enough here too; final output ordering happens once + // after selectedCandidates is complete. + std::nth_element(remainingCandidates.begin(), + remainingCandidates.begin() + numFillCandidates, + remainingCandidates.end(), + higherAttributedCandidate); + } + + for (int32_t candidateIndex = 0; candidateIndex < numFillCandidates; ++candidateIndex) { + selectedCandidates.push_back(remainingCandidates[candidateIndex]); + } + } + + // The target/fill passes decide membership; final output order is still by + // best calibrated score so downstream sequence construction sees a ranked + // PPR sequence rather than channel-grouped blocks. + if (selectedCandidates.size() > 1) { + std::sort(selectedCandidates.begin(), selectedCandidates.end(), higherAttributedCandidate); + } + + std::vector selectedNodes; + selectedNodes.reserve(selectedCandidates.size()); + for (const auto& candidate : selectedCandidates) { + selectedNodes.push_back(candidate.nodeId); + } + return selectedNodes; +} + +PPRExtractResult PPRForwardPush::extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) { TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, "."); - std::unordered_map> extractedPPRByNodeTypeId; + PPRExtractResult extractedPPRByNodeTypeId; // Emit an entry for every node type, even if unreachable in this batch (empty tensors, // all-zero valid_counts). This keeps the output shape consistent across batches so // downstream model architectures see a fixed set of PPR edge types every iteration. @@ -420,6 +700,97 @@ std::unordered_map& states, + const std::vector& channelTargetCounts, + int32_t maxPPRNodes, + bool enableResidualTopUp) { + // Typed channels are constructed from the same seed batch and graph schema; + // only the edge-type traversal allowlist differs. The sampler calls typed + // extraction only when typed-channel target counts are configured, so at + // least one state exists. Use the first state as the shared schema source + // for batch size and node-type count. + const auto* firstState = states.front(); + int32_t batchSize = firstState->_batchSize; + int32_t numNodeTypes = firstState->_numNodeTypes; + int32_t numChannels = static_cast(states.size()); + // Feature width is 1 + 2C: + // column 0: best calibrated score across channels, used as the scalar + // PPR weight by downstream ranking/sequence construction. + // columns [1, C]: calibrated score for each channel. + // columns [1 + C, 1 + 2C): presence bit for each channel. + int32_t numEdgeAttrFeatures = 1 + (2 * numChannels); + + // Pre-size the per-seed feature map below. Each channel can contribute up + // to maxPPRNodes candidates before cross-channel dedup and target filling. + size_t outputCandidateReserveSize = static_cast(numChannels) * static_cast(maxPPRNodes); + + PPRExtractResult extractedTypedPPRByNodeTypeId; + + for (int32_t nodeTypeId = 0; nodeTypeId < numNodeTypes; ++nodeTypeId) { + std::vector flatIds; + std::vector flatFeatureValues; + std::vector validCounts; + + for (int32_t seedIdx = 0; seedIdx < batchSize; ++seedIdx) { + std::unordered_map> outputScores; + outputScores.reserve(outputCandidateReserveSize); + // outputCandidatesByChannel includes finalized PPR rows, plus + // residual-aware rows when top-up is enabled. Selection and emitted + // features use this same view so both finalized and residual + // candidates obey the per-channel target counts. + std::vector>> outputCandidatesByChannel( + static_cast(numChannels)); + + for (int32_t channelIndex = 0; channelIndex < numChannels; ++channelIndex) { + const auto& nodeTypeState = states[channelIndex]->_state[seedIdx][nodeTypeId]; + auto outputNodesAndScores = selectFinalizedPPRPairs(nodeTypeState, maxPPRNodes); + + if (enableResidualTopUp) { + addResidualMassToPPRPairs(nodeTypeState, outputNodesAndScores); + appendResidualTopUpPairs(nodeTypeState, outputNodesAndScores, maxPPRNodes); + } + + auto outputMaxScoreIter = + std::max_element(outputNodesAndScores.begin(), + outputNodesAndScores.end(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + double outputMaxScore = + outputMaxScoreIter != outputNodesAndScores.end() ? outputMaxScoreIter->second : 0.0; + + outputCandidatesByChannel[channelIndex].reserve(outputNodesAndScores.size()); + addTypedPPRSeedFeaturesAndCandidates(outputNodesAndScores, + outputMaxScore, + channelIndex, + numChannels, + outputScores, + outputCandidatesByChannel[channelIndex]); + } + + auto selectedNodes = selectTypedPPRNodeIds(outputCandidatesByChannel, channelTargetCounts, maxPPRNodes); + int32_t selectedNodeCount = static_cast(selectedNodes.size()); + + for (int32_t nodeId : selectedNodes) { + flatIds.push_back(static_cast(nodeId)); + const auto& features = outputScores.at(nodeId); + flatFeatureValues.insert(flatFeatureValues.end(), features.begin(), features.end()); + } + + validCounts.push_back(static_cast(selectedNodeCount)); + } + + auto flatWeights = + torch::tensor(flatFeatureValues, torch::kDouble) + .reshape({static_cast(flatIds.size()), static_cast(numEdgeAttrFeatures)}); + extractedTypedPPRByNodeTypeId[nodeTypeId] = { + torch::tensor(flatIds, torch::kLong), + flatWeights, + torch::tensor(validCounts, torch::kLong), + }; + } + + return extractedTypedPPRByNodeTypeId; +} + int32_t PPRForwardPush::getTotalDegree(int32_t nodeId, int32_t nodeTypeId) const { TORCH_CHECK(nodeTypeId >= 0, "nodeTypeId ", nodeTypeId, " is negative, which indicates a sampler bug."); TORCH_CHECK(nodeTypeId < static_cast(_degreeTensors.size()), diff --git a/gigl-core/core/sampling/ppr_forward_push.h b/gigl-core/core/sampling/ppr_forward_push.h index 4698a6998..e909fa479 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/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; +using NeighborFetchMap = std::unordered_map; + +// PPR extraction output, keyed by integer node type ID: +// ids, weights/edge_attr, valid_counts. +using PPRExtractTensors = std::tuple; +using PPRExtractResult = std::unordered_map; + // 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 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 drainedChannelIndices; + + // Subset of drainedChannelIndices that still have fetch budget and at least + // one non-empty uncached frontier. + std::vector fetchChannelIndices; + + // Edge types requested by each fetch channel, aligned with fetchChannelIndices. + std::vector> 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 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> 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>& 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> extractTopKWithResidualTopUp( - int32_t maxPPRNodes, bool enableResidualTopUp); + PPRExtractResult extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp); + + friend PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, + const std::vector& 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> _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& states, + const std::vector& 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& states, + const std::vector& channelTargetCounts, + int32_t maxPPRNodes, + bool enableResidualTopUp); + } // namespace gigl diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index b92b21418..9e1fd9f44 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/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> neighborTensorsByEtypeId; + NeighborFetchMap neighborTensorsByEtypeId; // Dict iteration touches Python objects — GIL must be held here. for (auto item : fetchedByEtypeId) { auto edgeTypeId = item.first.cast(); @@ -43,20 +43,21 @@ static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedB } static std::optional> drainQueueWrapper(PPRForwardPush& state) { - std::optional> drained; + std::optional> 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> -extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) { - std::unordered_map> 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& fetchIterationCounts, + int32_t maxFetchIterations) { + std::vector 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()); + } + + // 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& channelTargetCounts, + int32_t maxPPRNodes, + bool enableResidualTopUp) { + std::vector 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()); + } + + // 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")); } diff --git a/gigl-core/src/gigl_core/__init__.py b/gigl-core/src/gigl_core/__init__.py index 524135619..0ad38b6f2 100644 --- a/gigl-core/src/gigl_core/__init__.py +++ b/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", +] diff --git a/gigl-core/src/gigl_core/ppr_forward_push.pyi b/gigl-core/src/gigl_core/ppr_forward_push.pyi index 992315761..0ec80611a 100644 --- a/gigl-core/src/gigl_core/ppr_forward_push.pyi +++ b/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]]: ... diff --git a/gigl-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index b9f07125b..ceee2d6e6 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -2,7 +2,11 @@ #include "sampling/ppr_forward_push.h" #include +#include +using gigl::drainTypedPPRChannelQueues; +using gigl::extractTypedTopKWithResidualTopUp; +using gigl::NeighborFetchMap; using gigl::PPRForwardPush; // Builds a single-edge-type, single-node-type PPRForwardPush. @@ -21,17 +25,35 @@ static PPRForwardPush makeState(const std::vector& seeds, // Convenience wrapper: build the fetchedByEtypeId argument for pushResiduals // from flat vectors, keeping test call sites readable. -static std::unordered_map> makeFetched( - int32_t edgeTypeId, - const std::vector& nodeIds, - const std::vector& flatNeighborIds, - const std::vector& counts) { +static NeighborFetchMap makeFetched(int32_t edgeTypeId, + const std::vector& nodeIds, + const std::vector& flatNeighborIds, + const std::vector& counts) { return {{edgeTypeId, {torch::tensor(nodeIds, torch::kLong), torch::tensor(flatNeighborIds, torch::kLong), torch::tensor(counts, torch::kLong)}}}; } +static std::unordered_set tensorValues(const torch::Tensor& values) { + std::unordered_set result; + auto accessor = values.accessor(); + for (int64_t index = 0; index < values.size(0); ++index) { + result.insert(accessor[index]); + } + return result; +} + +static std::vector tensorToInt64Vector(const torch::Tensor& values) { + std::vector result; + auto accessor = values.accessor(); + result.reserve(static_cast(values.size(0))); + for (int64_t index = 0; index < values.size(0); ++index) { + result.push_back(accessor[index]); + } + return result; +} + // After construction, drainQueue() returns the seed node under etype 0. TEST(PPRForwardPush, DrainQueueReturnsSeedNodeInitially) { auto state = makeState(/*seeds=*/{0}, /*alpha=*/0.15, /*requeueThresholdFactor=*/1e-6, /*degrees=*/{1}); @@ -110,6 +132,171 @@ TEST(PPRForwardPush, NeighborCacheAvoidsRefetchingPreviouslyFetchedNode) { EXPECT_TRUE(iter3->empty()); } +TEST(PPRForwardPush, DrainTypedPPRChannelQueuesUnionsChannelFrontiers) { + auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1}); + auto channel1 = makeState(/*seeds=*/{1}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1}); + + auto queueDrainResult = drainTypedPPRChannelQueues(std::vector{&channel0, &channel1}, + /*fetchIterationCounts=*/{0, 0}, + /*maxFetchIterations=*/-1); + + EXPECT_EQ(queueDrainResult.drainedChannelIndices, std::vector({0, 1})); + EXPECT_EQ(queueDrainResult.fetchChannelIndices, std::vector({0, 1})); + ASSERT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel.size(), 2); + EXPECT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel[0], std::vector({0})); + EXPECT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel[1], std::vector({0})); + + ASSERT_NE(queueDrainResult.unionedNodeIdsByEdgeTypeId.find(0), queueDrainResult.unionedNodeIdsByEdgeTypeId.end()); + EXPECT_EQ(tensorValues(queueDrainResult.unionedNodeIdsByEdgeTypeId.at(0)), std::unordered_set({0, 1})); +} + +TEST(PPRForwardPush, DrainTypedPPRChannelQueuesHonorsPerChannelFetchBudget) { + auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1}); + auto channel1 = makeState(/*seeds=*/{1}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1}); + + auto queueDrainResult = drainTypedPPRChannelQueues(std::vector{&channel0, &channel1}, + /*fetchIterationCounts=*/{1, 0}, + /*maxFetchIterations=*/1); + + EXPECT_EQ(queueDrainResult.drainedChannelIndices, std::vector({0, 1})); + EXPECT_EQ(queueDrainResult.fetchChannelIndices, std::vector({1})); + ASSERT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel.size(), 1); + EXPECT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel[0], std::vector({0})); + + ASSERT_NE(queueDrainResult.unionedNodeIdsByEdgeTypeId.find(0), queueDrainResult.unionedNodeIdsByEdgeTypeId.end()); + EXPECT_EQ(tensorValues(queueDrainResult.unionedNodeIdsByEdgeTypeId.at(0)), std::unordered_set({1})); +} + +TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpMergesChannelsInCpp) { + std::vector degrees(21, 1); + auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + auto channel1 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + + channel0.drainQueue(); + channel0.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{10}, /*counts=*/{1})); + channel1.drainQueue(); + channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1})); + + auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel0, &channel1}, + /*channelTargetCounts=*/{2, 1}, + /*maxPPRNodes=*/3, + /*enableResidualTopUp=*/true); + + ASSERT_NE(result.find(0), result.end()); + const auto& [ids, features, counts] = result.at(0); + EXPECT_EQ(tensorToInt64Vector(ids), std::vector({0, 10, 20})); + EXPECT_EQ(tensorToInt64Vector(counts), std::vector({3})); + ASSERT_EQ(features.size(0), 3); + ASSERT_EQ(features.size(1), 5); + + auto featureAccessor = features.accessor(); + EXPECT_NEAR(featureAccessor[0][0], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][1], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][2], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][3], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][4], 1.0, 1e-9); + + EXPECT_NEAR(featureAccessor[1][0], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[1][1], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[1][2], 0.0, 1e-9); + EXPECT_NEAR(featureAccessor[1][3], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[1][4], 0.0, 1e-9); + + EXPECT_NEAR(featureAccessor[2][0], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[2][1], 0.0, 1e-9); + EXPECT_NEAR(featureAccessor[2][2], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[2][3], 0.0, 1e-9); + EXPECT_NEAR(featureAccessor[2][4], 1.0, 1e-9); +} + +TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpUsesTargetsForResidualRows) { + std::vector degrees(21, 1); + auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + auto channel1 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + + channel0.drainQueue(); + channel0.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{10}, /*counts=*/{1})); + channel1.drainQueue(); + channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1})); + + auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel0, &channel1}, + /*channelTargetCounts=*/{1, 1}, + /*maxPPRNodes=*/2, + /*enableResidualTopUp=*/true); + + ASSERT_NE(result.find(0), result.end()); + const auto& [ids, features, counts] = result.at(0); + EXPECT_EQ(tensorToInt64Vector(ids), std::vector({0, 20})); + EXPECT_EQ(tensorToInt64Vector(counts), std::vector({2})); + ASSERT_EQ(features.size(0), 2); + ASSERT_EQ(features.size(1), 5); + + auto featureAccessor = features.accessor(); + EXPECT_NEAR(featureAccessor[0][0], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][1], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][2], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][3], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][4], 1.0, 1e-9); + + EXPECT_NEAR(featureAccessor[1][0], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[1][1], 0.0, 1e-9); + EXPECT_NEAR(featureAccessor[1][2], 0.5, 1e-9); + EXPECT_NEAR(featureAccessor[1][3], 0.0, 1e-9); + EXPECT_NEAR(featureAccessor[1][4], 1.0, 1e-9); +} + +TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpEmitsResidualAwareBaseRows) { + auto channel = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1e-9, /*degrees=*/{1, 1}); + + channel.drainQueue(); + channel.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{1}, /*counts=*/{1})); + channel.drainQueue(); + channel.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{1}, /*flatNeighborIds=*/{0}, /*counts=*/{1})); + + auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel}, + /*channelTargetCounts=*/{2}, + /*maxPPRNodes=*/2, + /*enableResidualTopUp=*/true); + + ASSERT_NE(result.find(0), result.end()); + const auto& [ids, features, counts] = result.at(0); + EXPECT_EQ(tensorToInt64Vector(ids), std::vector({0, 1})); + EXPECT_EQ(tensorToInt64Vector(counts), std::vector({2})); + ASSERT_EQ(features.size(0), 2); + ASSERT_EQ(features.size(1), 3); + + auto featureAccessor = features.accessor(); + EXPECT_NEAR(featureAccessor[0][0], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][1], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[0][2], 1.0, 1e-9); + EXPECT_NEAR(featureAccessor[1][0], 0.4, 1e-9); + EXPECT_NEAR(featureAccessor[1][1], 0.4, 1e-9); + EXPECT_NEAR(featureAccessor[1][2], 1.0, 1e-9); +} + +TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpCanDisableTopUp) { + std::vector degrees(21, 1); + auto channel0 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + auto channel1 = makeState(/*seeds=*/{0}, /*alpha=*/0.5, /*requeueThresholdFactor=*/1.0, degrees); + + channel0.drainQueue(); + channel0.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{10}, /*counts=*/{1})); + channel1.drainQueue(); + channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1})); + + auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel0, &channel1}, + /*channelTargetCounts=*/{2, 1}, + /*maxPPRNodes=*/3, + /*enableResidualTopUp=*/false); + + ASSERT_NE(result.find(0), result.end()); + const auto& [ids, features, counts] = result.at(0); + EXPECT_EQ(tensorToInt64Vector(ids), std::vector({0})); + EXPECT_EQ(tensorToInt64Vector(counts), std::vector({1})); + ASSERT_EQ(features.size(0), 1); + ASSERT_EQ(features.size(1), 5); +} + // Two seeds (0 and 1) both push residual to sink node 2. The neighbor-lookup // request must deduplicate to one entry for node 2, yet both seeds must still // accumulate a PPR score for it. diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 78de00436..bcffb5ff0 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -6,7 +6,11 @@ # TODO: Once gigl_core has a stable Python interface, re-export PPRForwardPush # under a gigl.core namespace rather than importing directly from the C++ extension. -from gigl_core import PPRForwardPush +from gigl_core import ( + PPRForwardPush, + drain_typed_ppr_channel_queues, + extract_typed_top_k_with_residual_top_up, +) from graphlearn_torch.sampler import ( HeteroSamplerOutput, NeighborOutput, @@ -59,6 +63,17 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): the PPR algorithm traverses across all edge types, switching edge types based on the current node type and the configured edge direction. + Internal execution follows the same shape for regular and typed PPR. Regular + PPR owns one C++ ``PPRForwardPush`` state per seed type: ``drain_queue`` + exposes the next frontier, Python performs the distributed neighbor fetch, + ``push_residuals`` updates the state, and C++ extraction emits the final + top-k plus residual top-up output. Typed PPR runs one ``PPRForwardPush`` + state per traversal channel. Its typed drain step unions the channel + frontiers before the same distributed fetch, pushes results back into the + active channel states, and then uses one typed C++ extraction step to apply + channel target counts, deduplicate shared candidates, and emit the final + typed edge-attribute features. + The ``edge_index`` and ``edge_attr`` fields on the output Data/HeteroData objects are populated with PPR seed-to-neighbor relationships (not edges in the original graph). ``N`` is the total number of (seed, neighbor) @@ -147,6 +162,8 @@ def __init__( ] self._is_homogeneous = True + self._typed_ppr_channel_target_counts: Optional[list[int]] = None + # Convert the public homogeneous/heterogeneous degree-tensor shape to # the node-type keyed form used internally by PPR. self._node_type_to_total_degree = self._convert_degree_tensors_to_dict( @@ -209,6 +226,10 @@ 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]] + ] = [] + def _convert_degree_tensors_to_dict( self, degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], @@ -304,7 +325,6 @@ def _extract_ppr_state_top_k( self, ppr_state, device: torch.device, - max_ppr_nodes: int, ) -> tuple[ Union[torch.Tensor, dict[NodeType, torch.Tensor]], Union[torch.Tensor, dict[NodeType, torch.Tensor]], @@ -334,7 +354,7 @@ def _extract_ppr_state_top_k( node_type_to_valid_counts: dict[NodeType, torch.Tensor] = {} extracted_results = ppr_state.extract_top_k_with_residual_top_up( - max_ppr_nodes, + self._max_ppr_nodes, self._enable_residual_topup, ) @@ -351,10 +371,6 @@ def _extract_ppr_state_top_k( node_type_to_valid_counts[node_type] = valid_counts.to(device) if self._is_homogeneous: - assert ( - len(node_type_to_flat_ids) == 1 - and DEFAULT_HOMOGENEOUS_NODE_TYPE in node_type_to_flat_ids - ) return ( node_type_to_flat_ids[DEFAULT_HOMOGENEOUS_NODE_TYPE], node_type_to_flat_weights[DEFAULT_HOMOGENEOUS_NODE_TYPE], @@ -367,6 +383,41 @@ def _extract_ppr_state_top_k( node_type_to_valid_counts, ) + def _extract_typed_ppr_state_top_k( + self, + ppr_states, + typed_ppr_channel_target_counts: list[int], + device: torch.device, + ) -> tuple[ + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + ]: + """Extract typed PPR results and move output tensors to the sampler device.""" + extracted_results = extract_typed_top_k_with_residual_top_up( + ppr_states, + typed_ppr_channel_target_counts, + self._max_ppr_nodes, + self._enable_residual_topup, + ) + node_type_to_flat_ids: dict[NodeType, torch.Tensor] = {} + node_type_to_flat_weights: dict[NodeType, torch.Tensor] = {} + node_type_to_valid_counts: dict[NodeType, torch.Tensor] = {} + for node_type_id, ( + flat_ids, + flat_weights, + valid_counts, + ) in extracted_results.items(): + node_type = self._ntype_id_to_ntype[node_type_id] + node_type_to_flat_ids[node_type] = flat_ids.to(device) + node_type_to_flat_weights[node_type] = flat_weights.to(device) + node_type_to_valid_counts[node_type] = valid_counts.to(device) + return ( + node_type_to_flat_ids, + node_type_to_flat_weights, + node_type_to_valid_counts, + ) + async def _compute_ppr_scores( self, seed_nodes: torch.Tensor, @@ -473,20 +524,131 @@ async def _compute_ppr_scores( # Fetch budget exhausted; push_residuals will use the existing neighbor cache. fetched_by_edge_type_id = {} - # Run in executor so the C++ push doesn't block the asyncio event loop. await loop.run_in_executor( - None, ppr_state.push_residuals, fetched_by_edge_type_id + None, + ppr_state.push_residuals, + fetched_by_edge_type_id, ) - # Regular sampling uses residual top-up as fill-only under the - # max_ppr_nodes cap. If finalized PPR scores already fill that cap, - # there is no remaining budget for residual candidates. return await loop.run_in_executor( None, self._extract_ppr_state_top_k, ppr_state, device, - self._max_ppr_nodes, + ) + + async def _compute_typed_ppr_scores( + self, + seed_nodes: torch.Tensor, + seed_node_type: NodeType, + typed_ppr_channel_target_counts: list[int], + ) -> tuple[ + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + ]: + """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++ deduplicates + candidates, attributes each node to its strongest channel, fills channel + target counts, redistributes unused slots by score, and emits typed + edge-attribute features 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_target_counts: Per-channel target output counts, 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. + """ + device = seed_nodes.device + loop = asyncio.get_running_loop() + + # Build one Forward Push state per typed channel. All states use the + # same seeds, restart probability, degree tensors, and destination-type + # map; only the per-node-type edge traversal allowlist differs. + def build_ppr_states() -> list[PPRForwardPush]: + return [ + PPRForwardPush( + seed_nodes, + self._node_type_to_id[seed_node_type], + self._alpha, + self._requeue_threshold_factor, + node_type_id_to_edge_type_ids, + self._edge_type_id_to_dst_ntype_id, + self._degree_tensors_for_cpp, + ) + for node_type_id_to_edge_type_ids in ( + self._typed_ppr_channel_to_node_type_id_to_edge_type_ids + ) + ] + + ppr_states = await loop.run_in_executor(None, build_ppr_states) + fetch_iteration_counts = [0 for _ in ppr_states] + max_fetch_iterations = ( + self._max_fetch_iterations if self._max_fetch_iterations is not None else -1 + ) + + while True: + ( + drained_channel_indices, + fetch_channel_indices, + edge_type_ids_by_fetch_channel, + unioned_node_ids_by_edge_type_id, + ) = await loop.run_in_executor( + None, + drain_typed_ppr_channel_queues, + ppr_states, + fetch_iteration_counts, + max_fetch_iterations, + ) + if not drained_channel_indices: + break + + fetched_by_channel: list[ + dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + ] = [dict() for _ in ppr_states] + + if unioned_node_ids_by_edge_type_id: + union_fetched_by_edge_type_id = await self._batch_fetch_neighbors( + unioned_node_ids_by_edge_type_id + ) + for channel_index, edge_type_ids in zip( + fetch_channel_indices, + edge_type_ids_by_fetch_channel, + strict=True, + ): + fetch_iteration_counts[channel_index] += 1 + fetched_by_channel[channel_index] = { + edge_type_id: union_fetched_by_edge_type_id[edge_type_id] + for edge_type_id in edge_type_ids + } + + # Push every non-converged channel. The fetched_by_channel entry is + # empty for channels that have no new fetch work; PPRForwardPush will + # use its cached neighbors in that case. + push_tasks = [ + loop.run_in_executor( + None, + ppr_states[channel_index].push_residuals, + fetched_by_channel[channel_index], + ) + for channel_index in drained_channel_indices + ] + await asyncio.gather(*push_tasks) + + return await loop.run_in_executor( + None, + self._extract_typed_ppr_state_top_k, + ppr_states, + typed_ppr_channel_target_counts, + device, ) async def _sample_from_nodes( @@ -590,15 +752,28 @@ async def _sample_from_nodes( # which is most beneficial when there are 2+ distinct seed node types # (e.g. cross-type supervision edges like user→story). seed_types = list(nodes_by_seed_type.keys()) - ppr_results = await asyncio.gather( - *[ - self._compute_ppr_scores( - nodes_by_seed_type[seed_type], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. - seed_type, - ) - for seed_type in seed_types - ] - ) + typed_ppr_channel_target_counts = self._typed_ppr_channel_target_counts + if typed_ppr_channel_target_counts is None: + ppr_results = await asyncio.gather( + *[ + self._compute_ppr_scores( + nodes_by_seed_type[seed_type], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + seed_type, + ) + for seed_type in seed_types + ] + ) + else: + ppr_results = await asyncio.gather( + *[ + self._compute_typed_ppr_scores( + nodes_by_seed_type[seed_type], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + seed_type, + typed_ppr_channel_target_counts, + ) + for seed_type in seed_types + ] + ) neighbor_dict: dict[EdgeType, list[torch.Tensor]] = {} ppr_edge_type_to_flat_weights: dict[EdgeType, torch.Tensor] = {} @@ -614,16 +789,16 @@ async def _sample_from_nodes( for node_type, flat_ids in node_type_to_flat_ids.items(): ppr_edge_type: EdgeType = (seed_type, "ppr", node_type) - valid_counts = node_type_to_valid_counts[node_type] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + valid_counts = node_type_to_valid_counts[node_type] ppr_edge_type_to_flat_weights[ppr_edge_type] = ( - node_type_to_flat_weights[node_type] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + node_type_to_flat_weights[node_type] ) # Skip empty pairs; induce_next handles deduplication across # seed types so a neighbor reachable from multiple seed types # gets one consistent local index in node_dict[node_type]. - if flat_ids.numel() > 0: # ty: ignore[unresolved-attribute] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. - neighbor_dict[ppr_edge_type] = [ # ty: ignore[invalid-assignment] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + if flat_ids.numel() > 0: + neighbor_dict[ppr_edge_type] = [ source_dict[seed_type], flat_ids, valid_counts, @@ -669,9 +844,7 @@ async def _sample_from_nodes( edge_index = torch.stack([rows, cols]) else: edge_index = torch.zeros(2, 0, dtype=torch.long, device=self.device) - flat_weights = torch.zeros( - 0, dtype=torch.double, device=self.device - ) + flat_weights = flat_weights.new_zeros((0, *flat_weights.shape[1:])) edge_type_repr = repr(ppr_edge_type) metadata[f"{PPR_EDGE_INDEX_METADATA_KEY}{edge_type_repr}"] = edge_index metadata[f"{PPR_WEIGHT_METADATA_KEY}{edge_type_repr}"] = flat_weights