From 41b08d4053f8bff08eaf941991a476c19ec36ad1 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 16:04:09 +0000 Subject: [PATCH 01/10] Add typed PPR C++ extraction support --- gigl-core/core/sampling/ppr_forward_push.cpp | 396 +++++++++++++++++++ gigl-core/core/sampling/ppr_forward_push.h | 84 ++++ gigl-core/tests/ppr_forward_push_test.cpp | 153 +++++++ 3 files changed, 633 insertions(+) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 915df7c56..10e9454f8 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -147,6 +147,60 @@ std::optional> PPRForwardPush::drainQ return result; } +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; + std::unordered_map> unionedFrontierNodeIdsByEdgeTypeId; + + // TODO: If typed queue draining becomes a measured blocker, evaluate + // parallelizing this channel loop with per-thread frontier maps and a + // deterministic merge into queueDrainResult. + for (size_t channelIndex = 0; channelIndex < states.size(); ++channelIndex) { + PPRForwardPush* state = states[channelIndex]; + + auto channelFrontierByEdgeTypeId = state->drainQueue(); + 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& unionedFrontierNodeIds = unionedFrontierNodeIdsByEdgeTypeId[edgeTypeId]; + for (int64_t nodeIndex = 0; nodeIndex < nodes.size(0); ++nodeIndex) { + unionedFrontierNodeIds.insert(nodeAccessor[nodeIndex]); + } + } + + if (!requestedEdgeTypeIds.empty()) { + queueDrainResult.fetchChannelIndices.push_back(static_cast(channelIndex)); + queueDrainResult.edgeTypeIdsByFetchChannel.push_back(std::move(requestedEdgeTypeIds)); + } + } + + for (const auto& [edgeTypeId, unionedFrontierNodeIds] : unionedFrontierNodeIdsByEdgeTypeId) { + std::vector nodeIdsToLookup(unionedFrontierNodeIds.begin(), unionedFrontierNodeIds.end()); + queueDrainResult.unionedNodeIdsByEdgeTypeId[edgeTypeId] = torch::tensor(nodeIdsToLookup, torch::kLong); + } + return queueDrainResult; +} + void PPRForwardPush::pushResiduals( const std::unordered_map>& fetchedByEtypeId) { // Step 1: Persist fetched neighbor lists in the per-state cache. drainQueue() @@ -374,6 +428,211 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, } } +// Helper function for adding one channel's extracted PPR candidates into a +// selection-only candidate list. +// +// This intentionally does not write edge_attr features. Typed extraction uses +// it for the base quota-biased pass, where residual top-up must not affect which +// nodes consume per-channel quota slots. +// +// Inputs: +// channelSelectionCandidates: mutable sortable candidates for this channel quota. +// nodesAndScores: extracted (node_id, score) candidates for one channel. +// maxScore: largest score in the channel, used to calibrate scores to [0, 1]. +// +// Expected output: channelSelectionCandidates contains this channel's calibrated +// candidates for later quota selection. +static void addTypedPPRChannelCandidates(std::vector>& channelSelectionCandidates, + const std::vector>& nodesAndScores, + double maxScore) { + for (const auto& [nodeId, score] : nodesAndScores) { + double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; + channelSelectionCandidates.emplace_back(nodeId, calibratedScore); + } +} + +// Helper function for adding one channel's extracted PPR candidates into the +// emitted typed feature table and a fill-pass candidate list. +// +// Unlike addTypedPPRChannelCandidates, this helper owns the output view: it +// populates the edge_attr row for every node it sees. When residual top-up is +// enabled, callers pass residual-aware scores here so base-selected and top-up +// rows are emitted on the same ppr_score + residual scale. +// +// Inputs: +// outputScoresByNodeId: mutable map from node ID to emitted edge_attr features. +// channelOutputCandidates: mutable sortable candidates for the fill pass. +// 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. +// +// Expected output: outputScoresByNodeId has this channel's score/presence +// features merged in, and channelOutputCandidates contains this channel's +// calibrated candidates. +static void addTypedPPRSeedFeaturesAndCandidates(std::unordered_map>& outputScoresByNodeId, + std::vector>& channelOutputCandidates, + const std::vector>& nodesAndScores, + double maxScore, + int32_t channelIndex, + int32_t numChannels) { + // Typed edge_attr rows store one best score across channels, followed by + // one score and one presence bit per channel. + int32_t numEdgeAttrFeatures = 1 + (2 * numChannels); + for (const auto& [nodeId, score] : nodesAndScores) { + double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; + 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. A node may be seen multiple times in the + // same channel view, so keep the strongest calibrated score. + scoreFeatures[channelScoreIndex] = std::max(scoreFeatures[channelScoreIndex], calibratedScore); + scoreFeatures[channelPresenceIndex] = 1.0; + + // Keep a per-channel sortable candidate list so channel quotas can be + // applied before the global cross-channel dedup pass. + channelOutputCandidates.emplace_back(nodeId, calibratedScore); + } +} + +// Helper function for applying typed channel quotas and cross-channel dedup. +// +// Inputs: +// candidatesByChannel: mutable (node_id, calibrated_score) candidates per channel. +// channelQuotas: maximum number of candidates to consider from each channel. +// maxPPRNodes: maximum number of deduplicated node IDs to return for a seed. +// +// Expected output: node IDs selected after per-channel quota filtering, then +// globally ranked by best calibrated score. Tie breakers keep output deterministic. +static std::vector selectTypedPPRNodeIds( + std::vector>>& candidatesByChannel, + const std::vector& channelQuotas, + int32_t maxPPRNodes) { + struct GlobalCandidate { + double bestCalibratedScore; + double channelCalibratedScore; + int32_t channelIndex; + int32_t nodeId; + }; + + // Per-channel ordering is by the channel-local score. Global ordering is by + // the node's best score across all channels, then by the candidate's own + // channel score and stable IDs for deterministic ties. + const auto higherScorePair = [](const auto& a, const auto& b) { + if (a.second != b.second) { + return a.second > b.second; + } + return a.first < b.first; + }; + const auto higherGlobalCandidate = [](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; + }; + + // Bound reserve sizes by the number of rows that can survive per-channel + // quotas. Current callers already cap channels to these limits, but this + // keeps the helper efficient if a future caller passes larger vectors. + size_t candidateRowReserveSize = 0; + for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByChannel.size()); ++channelIndex) { + candidateRowReserveSize += static_cast( + std::min(channelQuotas[channelIndex], static_cast(candidatesByChannel[channelIndex].size()))); + } + + // First apply per-channel quotas. candidateRows may still contain the same + // node multiple times if multiple channels reached it, while + // bestCalibratedScoreByNodeId tracks the cross-channel score used for final + // ranking. + std::unordered_map bestCalibratedScoreByNodeId; + bestCalibratedScoreByNodeId.reserve(candidateRowReserveSize); + std::vector candidateRows; + candidateRows.reserve(candidateRowReserveSize); + for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByChannel.size()); ++channelIndex) { + auto& candidates = candidatesByChannel[channelIndex]; + int32_t channelQuota = channelQuotas[channelIndex]; + int32_t numCandidates = std::min(channelQuota, static_cast(candidates.size())); + if (numCandidates <= 0) { + continue; + } + + // Only partition when a channel has more rows than its quota. We do not + // need a sorted per-channel prefix; we only need the top quota rows + // before the global cross-channel rank. + if (numCandidates < static_cast(candidates.size())) { + std::nth_element(candidates.begin(), candidates.begin() + numCandidates, candidates.end(), higherScorePair); + } + + for (int32_t candidateIndex = 0; candidateIndex < numCandidates; ++candidateIndex) { + int32_t nodeId = candidates[candidateIndex].first; + double calibratedScore = candidates[candidateIndex].second; + candidateRows.push_back({0.0, calibratedScore, channelIndex, nodeId}); + auto [bestScoreIter, inserted] = bestCalibratedScoreByNodeId.emplace(nodeId, calibratedScore); + if (!inserted) { + bestScoreIter->second = std::max(bestScoreIter->second, calibratedScore); + } + } + } + + // Collapse duplicate node IDs before global ranking. If the same node is + // present through multiple channels, keep the row that would sort highest + // under the final comparator; its bestCalibratedScore is still the best + // score observed for that node across every quota-surviving channel. + std::unordered_map bestCandidateByNodeId; + bestCandidateByNodeId.reserve(bestCalibratedScoreByNodeId.size()); + for (auto candidate : candidateRows) { + candidate.bestCalibratedScore = bestCalibratedScoreByNodeId.at(candidate.nodeId); + auto [candidateIter, inserted] = bestCandidateByNodeId.emplace(candidate.nodeId, candidate); + if (!inserted && higherGlobalCandidate(candidate, candidateIter->second)) { + candidateIter->second = candidate; + } + } + + // Move unique candidates into a vector so the final top-k ranking can use + // partial_sort. This avoids sorting the full candidate set when + // maxPPRNodes is smaller than the number of unique channel candidates. + std::vector globalCandidates; + globalCandidates.reserve(bestCandidateByNodeId.size()); + for (const auto& candidateEntry : bestCandidateByNodeId) { + globalCandidates.push_back(candidateEntry.second); + } + + int32_t selectedReserveSize = std::min(maxPPRNodes, static_cast(globalCandidates.size())); + if (selectedReserveSize < static_cast(globalCandidates.size())) { + std::partial_sort(globalCandidates.begin(), + globalCandidates.begin() + selectedReserveSize, + globalCandidates.end(), + higherGlobalCandidate); + } else if (globalCandidates.size() > 1) { + std::sort(globalCandidates.begin(), globalCandidates.end(), higherGlobalCandidate); + } + + // globalCandidates is already ranked through selectedReserveSize, so emit + // only node IDs in final order. Dedup already happened above. + std::vector selectedNodes; + selectedNodes.reserve(static_cast(selectedReserveSize)); + for (int32_t candidateIndex = 0; candidateIndex < selectedReserveSize; ++candidateIndex) { + selectedNodes.push_back(globalCandidates[candidateIndex].nodeId); + } + return selectedNodes; +} + std::unordered_map> PPRForwardPush:: extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) { TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, "."); @@ -420,6 +679,143 @@ std::unordered_map> extractTypedTopKWithResidualTopUp( + const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp) { + const auto* firstState = states.front(); + int32_t batchSize = firstState->_batchSize; + int32_t numNodeTypes = firstState->_numNodeTypes; + int32_t numChannels = static_cast(states.size()); + // Typed edge_attr rows store one best score across channels, followed by + // one score and one presence bit per channel. + int32_t numEdgeAttrFeatures = 1 + (2 * numChannels); + + // Pre-size the per-seed feature map below. Output features may receive up + // to maxPPRNodes candidates per channel when residual top-up is on. + size_t outputCandidateReserveSize = 0; + for (int32_t channelIndex = 0; channelIndex < numChannels; ++channelIndex) { + int32_t channelPPRNodeBudget = std::min(channelQuotas[channelIndex], maxPPRNodes); + outputCandidateReserveSize += static_cast(enableResidualTopUp ? maxPPRNodes : channelPPRNodeBudget); + } + + std::unordered_map> extractedTypedPPRByNodeTypeId; + std::vector topUpChannelQuotas(static_cast(numChannels), maxPPRNodes); + + 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); + // Keep the two typed views separate: + // - baseCandidatesByChannel is selection-only. It contains raw + // finalized PPR rows and drives the quota-biased first pass. + // - outputScores/outputCandidatesByChannel is the emitted view. + // It includes residual-aware scores when top-up is enabled, so + // base-selected and fill-selected rows share one score scale. + std::vector>> baseCandidatesByChannel( + static_cast(numChannels)); + std::vector>> outputCandidatesByChannel( + static_cast(numChannels)); + + for (int32_t channelIndex = 0; channelIndex < numChannels; ++channelIndex) { + const auto& nodeTypeState = states[channelIndex]->_state[seedIdx][nodeTypeId]; + int32_t channelPPRNodeBudget = std::min(channelQuotas[channelIndex], maxPPRNodes); + auto baseNodesAndScores = selectFinalizedPPRPairs(nodeTypeState, channelPPRNodeBudget); + + // Keep residual top-up out of the quota-biased base pass. The + // output view starts from the same finalized candidates, then + // adds residual mass/candidates for emitted features and the + // later fill pass. + auto outputNodesAndScores = baseNodesAndScores; + if (enableResidualTopUp) { + addResidualMassToPPRPairs(nodeTypeState, outputNodesAndScores); + appendResidualTopUpPairs(nodeTypeState, outputNodesAndScores, maxPPRNodes); + } + + double baseMaxScore = 0.0; + for (const auto& nodeAndScore : baseNodesAndScores) { + baseMaxScore = std::max(baseMaxScore, nodeAndScore.second); + } + + double outputMaxScore = 0.0; + for (const auto& nodeAndScore : outputNodesAndScores) { + outputMaxScore = std::max(outputMaxScore, nodeAndScore.second); + } + + baseCandidatesByChannel[channelIndex].reserve(baseNodesAndScores.size()); + outputCandidatesByChannel[channelIndex].reserve(outputNodesAndScores.size()); + addTypedPPRChannelCandidates(baseCandidatesByChannel[channelIndex], baseNodesAndScores, baseMaxScore); + addTypedPPRSeedFeaturesAndCandidates(outputScores, + outputCandidatesByChannel[channelIndex], + outputNodesAndScores, + outputMaxScore, + channelIndex, + numChannels); + } + + auto selectedNodes = selectTypedPPRNodeIds(baseCandidatesByChannel, channelQuotas, maxPPRNodes); + int32_t selectedNodeCount = static_cast(selectedNodes.size()); + if (selectedNodeCount < maxPPRNodes) { + std::unordered_set selectedNodeIds(selectedNodes.begin(), selectedNodes.end()); + auto topUpSelectedNodes = + selectTypedPPRNodeIds(outputCandidatesByChannel, topUpChannelQuotas, maxPPRNodes); + for (int32_t nodeId : topUpSelectedNodes) { + if (selectedNodeCount >= maxPPRNodes) { + break; + } + if (selectedNodeIds.find(nodeId) != selectedNodeIds.end()) { + continue; + } + ++selectedNodeCount; + selectedNodeIds.insert(nodeId); + selectedNodes.push_back(nodeId); + } + } + + if (enableResidualTopUp && selectedNodes.size() > 1) { + std::vector> selectedNodesAndScores; + selectedNodesAndScores.reserve(selectedNodes.size()); + for (int32_t nodeId : selectedNodes) { + selectedNodesAndScores.emplace_back(nodeId, outputScores.at(nodeId)[0]); + } + std::stable_sort(selectedNodesAndScores.begin(), + selectedNodesAndScores.end(), + [](const auto& a, const auto& b) { return a.second > b.second; }); + + selectedNodes.clear(); + selectedNodes.reserve(selectedNodesAndScores.size()); + for (const auto& selectedNodeAndScore : selectedNodesAndScores) { + selectedNodes.push_back(selectedNodeAndScore.first); + } + } + + 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..876ef9889 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/gigl-core/core/sampling/ppr_forward_push.h @@ -25,6 +25,33 @@ 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. + // + // 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, keyed by + // integer edge type ID. Tensor values are int64 node IDs. + 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. // @@ -75,6 +102,12 @@ class PPRForwardPush { std::unordered_map> extractTopKWithResidualTopUp( int32_t maxPPRNodes, bool enableResidualTopUp); + friend std::unordered_map> + extractTypedTopKWithResidualTopUp(const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp); + private: // Total out-degree of a node across all edge types. Returns 0 for sink nodes. [[nodiscard]] int32_t getTotalDegree(int32_t nodeId, int32_t nodeTypeId) const; @@ -116,4 +149,55 @@ 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 each +// channel state, 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. +// +// 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 channel, finalized PPR candidates feed the quota-biased base pass. +// Residual top-up candidates are kept in a separate extended view and are used +// only after channel quotas and cross-channel dedup leave unfilled sequence +// slots. The typed merge calibrates scores within each channel, deduplicates +// candidates seen through multiple channels, and emits per-node-type tensors. +// +// Inputs: +// states: Completed PPRForwardPush states, one per typed channel. +// channelQuotas: Per-channel finalized-PPR quotas, aligned with states. +// maxPPRNodes: Maximum number of deduplicated nodes to return per seed. +// enableResidualTopUp: Whether residual candidates may fill unfilled slots. +// +// 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. +std::unordered_map> extractTypedTopKWithResidualTopUp( + const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp); + } // namespace gigl diff --git a/gigl-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index b9f07125b..398511d50 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -2,7 +2,10 @@ #include "sampling/ppr_forward_push.h" #include +#include +using gigl::drainTypedPPRChannelQueues; +using gigl::extractTypedTopKWithResidualTopUp; using gigl::PPRForwardPush; // Builds a single-edge-type, single-node-type PPRForwardPush. @@ -32,6 +35,25 @@ static std::unordered_map 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,137 @@ 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}, + /*channelQuotas=*/{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, 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}, + /*channelQuotas=*/{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}, + /*channelQuotas=*/{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. From a9d9cb26e48f73ab3d1fdd637c9dd626c575927b Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 17:32:01 +0000 Subject: [PATCH 02/10] Clarify typed PPR C++ helper contracts --- gigl-core/core/sampling/ppr_forward_push.cpp | 99 +++++++++++-------- gigl-core/core/sampling/ppr_forward_push.h | 49 +++++---- .../core/sampling/python_ppr_forward_push.cpp | 9 +- gigl-core/tests/ppr_forward_push_test.cpp | 16 ++- 4 files changed, 101 insertions(+), 72 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 10e9454f8..5dd89800a 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -158,11 +158,15 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector> unionedFrontierNodeIdsByEdgeTypeId; - - // TODO: If typed queue draining becomes a measured blocker, evaluate - // parallelizing this channel loop with per-thread frontier maps and a - // deterministic merge into 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; + + // TODO: If benchmarking shows typed PPR is materially slower than regular + // heterogeneous PPR because of this drain loop, evaluate parallelizing + // channels with per-thread frontier maps and a deterministic merge into + // queueDrainResult. for (size_t channelIndex = 0; channelIndex < states.size(); ++channelIndex) { PPRForwardPush* state = states[channelIndex]; @@ -182,9 +186,9 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector(); - auto& unionedFrontierNodeIds = unionedFrontierNodeIdsByEdgeTypeId[edgeTypeId]; + auto& unionedSourceNodeIds = unionedSourceNodeIdsByEdgeTypeId[edgeTypeId]; for (int64_t nodeIndex = 0; nodeIndex < nodes.size(0); ++nodeIndex) { - unionedFrontierNodeIds.insert(nodeAccessor[nodeIndex]); + unionedSourceNodeIds.insert(nodeAccessor[nodeIndex]); } } @@ -194,15 +198,14 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector nodeIdsToLookup(unionedFrontierNodeIds.begin(), unionedFrontierNodeIds.end()); + 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 std::unordered_map>& fetchedByEtypeId) { +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 @@ -436,16 +439,19 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, // nodes consume per-channel quota slots. // // Inputs: -// channelSelectionCandidates: mutable sortable candidates for this channel quota. // nodesAndScores: extracted (node_id, score) candidates for one channel. // maxScore: largest score in the channel, used to calibrate scores to [0, 1]. +// channelSelectionCandidates: mutable sortable candidates for this channel quota. // // Expected output: channelSelectionCandidates contains this channel's calibrated // candidates for later quota selection. -static void addTypedPPRChannelCandidates(std::vector>& channelSelectionCandidates, - const std::vector>& nodesAndScores, - double maxScore) { +static void addTypedPPRChannelCandidates(const std::vector>& nodesAndScores, + double maxScore, + std::vector>& channelSelectionCandidates) { for (const auto& [nodeId, score] : nodesAndScores) { + // maxScore can be 0 only for an all-zero channel view. That is not + // expected for normal PPR mass, but keeping the guard avoids division + // by zero and makes empty/degenerate channel handling harmless. double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; channelSelectionCandidates.emplace_back(nodeId, calibratedScore); } @@ -460,26 +466,32 @@ static void addTypedPPRChannelCandidates(std::vector> // rows are emitted on the same ppr_score + residual scale. // // Inputs: -// outputScoresByNodeId: mutable map from node ID to emitted edge_attr features. -// channelOutputCandidates: mutable sortable candidates for the fill pass. // 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 the fill pass. // // Expected output: outputScoresByNodeId has this channel's score/presence // features merged in, and channelOutputCandidates contains this channel's // calibrated candidates. -static void addTypedPPRSeedFeaturesAndCandidates(std::unordered_map>& outputScoresByNodeId, - std::vector>& channelOutputCandidates, - const std::vector>& nodesAndScores, +static void addTypedPPRSeedFeaturesAndCandidates(const std::vector>& nodesAndScores, double maxScore, int32_t channelIndex, - int32_t numChannels) { - // Typed edge_attr rows store one best score across channels, followed by - // one score and one presence bit per channel. + int32_t numChannels, + std::unordered_map>& outputScoresByNodeId, + std::vector>& channelOutputCandidates) { + // 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) { + // maxScore can be 0 only for an all-zero channel view. That is not + // expected for normal PPR mass, but keeping the guard avoids division + // by zero and makes empty/degenerate channel handling harmless. double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; auto scoreIter = outputScoresByNodeId.find(nodeId); if (scoreIter == outputScoresByNodeId.end()) { @@ -615,11 +627,16 @@ static std::vector selectTypedPPRNodeIds( int32_t selectedReserveSize = std::min(maxPPRNodes, static_cast(globalCandidates.size())); if (selectedReserveSize < static_cast(globalCandidates.size())) { + // We only need the best maxPPRNodes rows, so partial_sort ranks the + // returned prefix without paying to fully sort candidates that will be + // discarded. std::partial_sort(globalCandidates.begin(), globalCandidates.begin() + selectedReserveSize, globalCandidates.end(), higherGlobalCandidate); } else if (globalCandidates.size() > 1) { + // If every unique candidate is returned, there is no discarded suffix; + // use a regular full sort to preserve the ordered output contract. std::sort(globalCandidates.begin(), globalCandidates.end(), higherGlobalCandidate); } @@ -633,11 +650,10 @@ static std::vector selectTypedPPRNodeIds( return selectedNodes; } -std::unordered_map> PPRForwardPush:: - extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) { +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. @@ -679,17 +695,22 @@ std::unordered_map> extractTypedTopKWithResidualTopUp( - const std::vector& states, - const std::vector& channelQuotas, - int32_t maxPPRNodes, - bool enableResidualTopUp) { +PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp) { + // Typed channels are constructed from the same seed batch and graph schema; + // only the edge-type traversal allowlist differs. 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()); - // Typed edge_attr rows store one best score across channels, followed by - // one score and one presence bit per channel. + // 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. Output features may receive up @@ -700,7 +721,7 @@ std::unordered_map(enableResidualTopUp ? maxPPRNodes : channelPPRNodeBudget); } - std::unordered_map> extractedTypedPPRByNodeTypeId; + PPRExtractResult extractedTypedPPRByNodeTypeId; std::vector topUpChannelQuotas(static_cast(numChannels), maxPPRNodes); for (int32_t nodeTypeId = 0; nodeTypeId < numNodeTypes; ++nodeTypeId) { @@ -749,13 +770,13 @@ std::unordered_map; +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) @@ -33,7 +43,10 @@ struct SeedNodeTypeState { // 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. + // Channels whose drainQueue() returned a value this iteration. Channel IDs + // are positional indices into the states/channelQuotas vectors. Python + // builds those vectors from typed_channel_quotas 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 @@ -47,8 +60,11 @@ struct TypedPPRQueueDrainResult { // Edge types requested by each fetch channel, aligned with fetchChannelIndices. std::vector> edgeTypeIdsByFetchChannel; - // Unioned node frontier for one shared distributed neighbor fetch, keyed by - // integer edge type ID. Tensor values are int64 node IDs. + // 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; }; @@ -78,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. // @@ -99,14 +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 std::unordered_map> - extractTypedTopKWithResidualTopUp(const std::vector& states, - const std::vector& channelQuotas, - int32_t maxPPRNodes, - bool enableResidualTopUp); + friend PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp); private: // Total out-degree of a node across all edge types. Returns 0 for sink nodes. @@ -194,10 +204,9 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector> extractTypedTopKWithResidualTopUp( - const std::vector& states, - const std::vector& channelQuotas, - int32_t maxPPRNodes, - bool enableResidualTopUp); +PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, + const std::vector& channelQuotas, + 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..35a715904 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(); @@ -54,9 +54,10 @@ static std::optional> drainQueueWrapp return drained; } -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. { diff --git a/gigl-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index 398511d50..b665b9ae1 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -6,6 +6,7 @@ using gigl::drainTypedPPRChannelQueues; using gigl::extractTypedTopKWithResidualTopUp; +using gigl::NeighborFetchMap; using gigl::PPRForwardPush; // Builds a single-edge-type, single-node-type PPRForwardPush. @@ -24,11 +25,10 @@ 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), @@ -146,8 +146,7 @@ TEST(PPRForwardPush, DrainTypedPPRChannelQueuesUnionsChannelFrontiers) { 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()); + ASSERT_NE(queueDrainResult.unionedNodeIdsByEdgeTypeId.find(0), queueDrainResult.unionedNodeIdsByEdgeTypeId.end()); EXPECT_EQ(tensorValues(queueDrainResult.unionedNodeIdsByEdgeTypeId.at(0)), std::unordered_set({0, 1})); } @@ -164,8 +163,7 @@ TEST(PPRForwardPush, DrainTypedPPRChannelQueuesHonorsPerChannelFetchBudget) { ASSERT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel.size(), 1); EXPECT_EQ(queueDrainResult.edgeTypeIdsByFetchChannel[0], std::vector({0})); - ASSERT_NE(queueDrainResult.unionedNodeIdsByEdgeTypeId.find(0), - queueDrainResult.unionedNodeIdsByEdgeTypeId.end()); + ASSERT_NE(queueDrainResult.unionedNodeIdsByEdgeTypeId.find(0), queueDrainResult.unionedNodeIdsByEdgeTypeId.end()); EXPECT_EQ(tensorValues(queueDrainResult.unionedNodeIdsByEdgeTypeId.at(0)), std::unordered_set({1})); } From 71c21ce4a484b980feaf42a4ff23357573c21a6d Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 18:43:59 +0000 Subject: [PATCH 03/10] Fail on invalid typed PPR channel max score --- gigl-core/core/sampling/ppr_forward_push.cpp | 25 +++++++++++++------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 5dd89800a..25dbc2c8a 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -448,11 +448,15 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, static void addTypedPPRChannelCandidates(const std::vector>& nodesAndScores, double maxScore, std::vector>& channelSelectionCandidates) { + if (nodesAndScores.empty()) { + return; + } + TORCH_CHECK(maxScore > 0.0, + "Typed PPR channel has candidates but non-positive max score ", + maxScore, + ", which indicates invalid PPR state."); for (const auto& [nodeId, score] : nodesAndScores) { - // maxScore can be 0 only for an all-zero channel view. That is not - // expected for normal PPR mass, but keeping the guard avoids division - // by zero and makes empty/degenerate channel handling harmless. - double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; + double calibratedScore = score / maxScore; channelSelectionCandidates.emplace_back(nodeId, calibratedScore); } } @@ -482,6 +486,14 @@ static void addTypedPPRSeedFeaturesAndCandidates(const std::vector>& 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. @@ -489,10 +501,7 @@ static void addTypedPPRSeedFeaturesAndCandidates(const std::vector 0.0 ? score / maxScore : 0.0; + double calibratedScore = score / maxScore; auto scoreIter = outputScoresByNodeId.find(nodeId); if (scoreIter == outputScoresByNodeId.end()) { scoreIter = outputScoresByNodeId.emplace(nodeId, std::vector(numEdgeAttrFeatures, 0.0)).first; From 7c2ac4f72d0d28435e59d5bbcd8c861f269fa58e Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 22:16:37 +0000 Subject: [PATCH 04/10] Move typed PPR sampler implementation into lower PR --- .../core/sampling/python_ppr_forward_push.cpp | 71 +++++- gigl-core/src/gigl_core/__init__.py | 12 +- gigl-core/src/gigl_core/ppr_forward_push.pyi | 14 ++ gigl/distributed/dist_ppr_sampler.py | 237 +++++++++++++++--- gigl/distributed/utils/dist_typed_sampler.py | 177 +++++++++++++ .../distributed/dist_typed_sampler_test.py | 77 ++++++ 6 files changed, 543 insertions(+), 45 deletions(-) create mode 100644 gigl/distributed/utils/dist_typed_sampler.py create mode 100644 tests/unit/distributed/dist_typed_sampler_test.py diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index 35a715904..ec1ca2d4a 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -43,15 +43,15 @@ 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 PPRExtractResult extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, @@ -67,6 +67,58 @@ static PPRExtractResult extractTopKWithResidualTopUpWrapper(PPRForwardPush& stat 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& channelQuotas, + 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, channelQuotas, maxPPRNodes, enableResidualTopUp); + } + return result; +} + } // namespace gigl // TORCH_EXTENSION_NAME is set by PyTorch's build system to match the Python @@ -86,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_quotas"), + 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..44bd82c45 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_quotas: Sequence[int], + max_ppr_nodes: int, + enable_residual_topup: bool, +) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: ... diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 78de00436..c14eb34c0 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, @@ -17,6 +21,11 @@ from graphlearn_torch.utils import merge_dict from gigl.distributed.base_sampler import BaseDistNeighborSampler +from gigl.distributed.utils.dist_typed_sampler import ( + HeteroPPRResult, + PPRResult, + TypedPPRChannelTraversalMaps, +) from gigl.types.graph import DEFAULT_HOMOGENEOUS_NODE_TYPE, is_label_edge_type # Trailing "." is an intentional separator. These constants are used both to @@ -59,6 +68,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 quotas, 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 +167,8 @@ def __init__( ] self._is_homogeneous = True + self._typed_ppr_channel_quotas: 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 +231,8 @@ def __init__( for node_type in all_node_types ] + self._typed_ppr_channel_to_node_type_id_to_edge_type_ids: TypedPPRChannelTraversalMaps = [] + def _convert_degree_tensors_to_dict( self, degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], @@ -304,12 +328,7 @@ 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]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - ]: + ) -> PPRResult: """Extract PPR neighbors from a completed C++ Forward Push state. The C++ kernel indexes node types by compact integer IDs for speed. @@ -334,7 +353,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 +370,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,15 +382,42 @@ def _extract_ppr_state_top_k( node_type_to_valid_counts, ) + def _extract_typed_ppr_state_top_k( + self, + ppr_states, + typed_ppr_channel_quotas: list[int], + device: torch.device, + ) -> HeteroPPRResult: + """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_quotas, + 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, seed_node_type: Optional[NodeType] = None, - ) -> tuple[ - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - ]: + ) -> PPRResult: """ Compute PPR scores for seed nodes using the push-based approximation algorithm. @@ -473,20 +515,126 @@ 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_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. + """ + 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_quotas, + device, ) async def _sample_from_nodes( @@ -590,15 +738,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_quotas = self._typed_ppr_channel_quotas + if typed_ppr_channel_quotas 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_quotas, + ) + 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 +775,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 +830,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 diff --git a/gigl/distributed/utils/dist_typed_sampler.py b/gigl/distributed/utils/dist_typed_sampler.py new file mode 100644 index 000000000..c836746a5 --- /dev/null +++ b/gigl/distributed/utils/dist_typed_sampler.py @@ -0,0 +1,177 @@ +"""Construction-time typed-PPR option parsing for distributed samplers. + +The helpers in this module validate typed-channel edge-type keys and convert +public edge-type keys into the compact integer traversal maps consumed by the +C++ forward-push kernel. They run once during ``DistPPRNeighborSampler`` +initialization, not in the per-batch PPR sampling hot loop. +""" + +from collections.abc import Sequence +from typing import Optional, Union, cast + +import torch +from graphlearn_torch.typing import EdgeType, NodeType + +# C++ PPR extraction output: flat node IDs, flat weights, and per-seed valid +# counts. Homogeneous extraction uses tensors directly; heterogeneous extraction +# uses dictionaries keyed by node type. +PPRResult = tuple[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], +] +# Heterogeneous-only view of PPRResult after typed PPR extraction. +HeteroPPRResult = tuple[ + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], +] +# Public typed_channel_quotas keys can be a single edge type or a grouped +# channel containing multiple edge types. +TypedPPRChannelKey = Union[EdgeType, tuple[EdgeType, ...]] + +"""TypedPPRChannelKey describes one public typed-PPR traversal channel key. + +A single canonical edge type creates one channel restricted to that edge type. +A tuple of canonical edge types creates one channel whose forward-push state may +traverse any edge type in the group. When typed PPR emits multi-column +``edge_attr`` tensors, channel columns follow the insertion order of the +``typed_channel_quotas`` mapping. +""" +# Parsed typed-channel edge-type allowlists, ordered to match the insertion +# order of typed_channel_quotas. +TypedPPRChannelEdgeTypeGroups = list[tuple[EdgeType, ...]] +# One channel's traversal map. The outer list is indexed by integer node-type +# ID; each inner list contains the integer edge-type IDs that channel may +# traverse from that node type. +TypedPPRChannelTraversalMap = list[list[int]] +# All typed-channel traversal maps, ordered to match typed-channel order. +TypedPPRChannelTraversalMaps = list[TypedPPRChannelTraversalMap] + + +def parse_typed_channel_quota_groups( + typed_channel_quotas: Optional[dict[TypedPPRChannelKey, int]], +) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[int]]]: + """Parse typed-PPR channel keys and split keys from quotas. + + Public options allow each channel key to be either one canonical edge type + or a non-empty tuple of canonical edge types. Internally, traversal setup + needs only the edge-type groups while merge selection needs the aligned + quota values, so this helper returns those two parallel lists. + + This is construction-time option parsing and is not part of the per-batch + PPR sampling hot loop. + + Args: + typed_channel_quotas: User-provided channel mapping from edge-type + allowlist to candidate quota. + + Returns: + ``(None, None)`` when typed PPR is disabled. Otherwise returns + ``(typed_channel_groups, typed_channel_quota_list)``, both ordered by + the input mapping insertion order. + + Raises: + ValueError: If a channel key is not a canonical edge type or non-empty + tuple of canonical edge types. + """ + if not typed_channel_quotas: + return None, None + + typed_channel_groups: TypedPPRChannelEdgeTypeGroups = [] + typed_channel_quota_list: list[int] = [] + + def is_canonical_edge_type(value: object) -> bool: + """Return whether ``value`` has PyG's canonical edge-type shape.""" + return ( + isinstance(value, tuple) + and len(value) == 3 + and all(isinstance(part, str) for part in value) + ) + + for edge_type_key, quota in typed_channel_quotas.items(): + if is_canonical_edge_type(edge_type_key): + edge_types = (cast(EdgeType, edge_type_key),) + elif ( + isinstance(edge_type_key, tuple) + and edge_type_key + and all(is_canonical_edge_type(edge_type) for edge_type in edge_type_key) + ): + edge_types = cast(tuple[EdgeType, ...], edge_type_key) + else: + raise ValueError( + "typed_channel_quotas keys must be a canonical edge type " + "(src_type, relation, dst_type) or a non-empty tuple of " + f"canonical edge types, got {edge_type_key!r}." + ) + typed_channel_groups.append(edge_types) + typed_channel_quota_list.append(quota) + + return typed_channel_groups, typed_channel_quota_list + + +def build_edge_type_channel_group_edge_type_ids( + edge_type_groups: TypedPPRChannelEdgeTypeGroups, + edge_type_to_edge_type_id: dict[EdgeType, int], + node_type_to_edge_types: dict[NodeType, list[EdgeType]], + node_types: Sequence[NodeType], +) -> TypedPPRChannelTraversalMaps: + """Convert typed-channel edge-type allowlists to PPRForwardPush IDs. + + Returns one traversal map per typed channel, ordered to match + ``edge_type_groups``. A single traversal map has shape + ``list[list[int]]``: the outer index is ``node_type_id``, and the inner list + contains allowed ``edge_type_id`` values for that node type in that channel. + + This conversion runs once during sampler construction; the resulting integer + maps are reused by the per-batch C++ PPR states. + + Args: + edge_type_groups: Ordered typed channels, where each channel is the + canonical edge types that its PPR state may traverse. + edge_type_to_edge_type_id: Mapping from canonical edge type to the + compact integer ID used by the C++ forward-push kernel. + node_type_to_edge_types: Traversable edge types keyed by anchor node + type, after label-edge filtering and edge-direction handling. + node_types: Ordered node types whose positions match the kernel's + integer node-type IDs. + + Returns: + ``channel_traversal_maps[channel_id][node_type_id]`` gives the allowed + integer edge-type IDs that channel may traverse from that node type. + + Raises: + ValueError: If a configured edge type is unknown, excluded from PPR + traversal, or cannot be traversed from any node type. + """ + known_edge_types = set(edge_type_to_edge_type_id.keys()) + channel_edge_type_ids_by_node_type: TypedPPRChannelTraversalMaps = [] + for channel_edge_types in edge_type_groups: + unknown_edge_types = set(channel_edge_types) - known_edge_types + if unknown_edge_types: + raise ValueError( + "typed_channel_quotas includes non-traversable edge types " + f"{sorted(unknown_edge_types)!r}. Known traversable edge " + f"types are {sorted(known_edge_types)!r}." + ) + + channel_edge_type_set = set(channel_edge_types) + node_type_id_to_channel_edge_type_ids: TypedPPRChannelTraversalMap = [] + for node_type in node_types: + channel_edge_type_ids_for_node_type: list[int] = [] + for edge_type in node_type_to_edge_types.get(node_type, []): + if edge_type in channel_edge_type_set: + channel_edge_type_ids_for_node_type.append( + edge_type_to_edge_type_id[edge_type] + ) + node_type_id_to_channel_edge_type_ids.append( + channel_edge_type_ids_for_node_type + ) + if not any(node_type_id_to_channel_edge_type_ids): + raise ValueError( + "typed_channel_quotas includes edge-type " + f"channel={channel_edge_types!r}, " + "but no traversable edge types exist for that channel." + ) + channel_edge_type_ids_by_node_type.append(node_type_id_to_channel_edge_type_ids) + return channel_edge_type_ids_by_node_type diff --git a/tests/unit/distributed/dist_typed_sampler_test.py b/tests/unit/distributed/dist_typed_sampler_test.py new file mode 100644 index 000000000..0d52f268a --- /dev/null +++ b/tests/unit/distributed/dist_typed_sampler_test.py @@ -0,0 +1,77 @@ +"""Unit tests for typed-PPR sampler construction helpers.""" + +from absl.testing import absltest + +from gigl.distributed.utils.dist_typed_sampler import ( + build_edge_type_channel_group_edge_type_ids, + parse_typed_channel_quota_groups, +) +from tests.test_assets.distributed.test_dataset import ( + STORY, + STORY_TO_USER, + USER, + USER_TO_STORY, +) +from tests.test_assets.test_case import TestCase + + +class DistTypedSamplerTest(TestCase): + def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( + self, + ) -> None: + """Verify typed-PPR can use canonical edge-type channels.""" + node_type_to_edge_types = { + USER: [USER_TO_STORY], + STORY: [STORY_TO_USER], + } + node_types = [USER, STORY] + edge_type_to_edge_type_id = { + USER_TO_STORY: 0, + STORY_TO_USER: 1, + } + + typed_channel_groups, typed_channel_quota_list = ( + parse_typed_channel_quota_groups( + { + USER_TO_STORY: 2, + (USER_TO_STORY, STORY_TO_USER): 3, + } + ) + ) + assert typed_channel_groups is not None + assert typed_channel_quota_list is not None + + self.assertEqual( + typed_channel_groups, + [ + (USER_TO_STORY,), + (USER_TO_STORY, STORY_TO_USER), + ], + ) + self.assertEqual(typed_channel_quota_list, [2, 3]) + self.assertEqual( + build_edge_type_channel_group_edge_type_ids( + edge_type_groups=typed_channel_groups, + edge_type_to_edge_type_id=edge_type_to_edge_type_id, + node_type_to_edge_types=node_type_to_edge_types, + node_types=node_types, + ), + [ + [[0], []], + [[0], [1]], + ], + ) + + with self.assertRaisesRegex(ValueError, "canonical edge type"): + parse_typed_channel_quota_groups({("bad",): 1}) + with self.assertRaisesRegex(ValueError, "non-traversable edge types"): + build_edge_type_channel_group_edge_type_ids( + edge_type_groups=[(("unknown", "edge", "type"),)], + edge_type_to_edge_type_id=edge_type_to_edge_type_id, + node_type_to_edge_types=node_type_to_edge_types, + node_types=node_types, + ) + + +if __name__ == "__main__": + absltest.main() From 72589824adf4e3f31ad4941c0392109f1441d74c Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 30 Jul 2026 06:45:36 +0000 Subject: [PATCH 05/10] Apply typed PPR channel target counts --- gigl-core/core/sampling/ppr_forward_push.cpp | 306 +++++++----------- gigl-core/core/sampling/ppr_forward_push.h | 27 +- .../core/sampling/python_ppr_forward_push.cpp | 6 +- gigl-core/src/gigl_core/ppr_forward_push.pyi | 2 +- gigl-core/tests/ppr_forward_push_test.cpp | 42 ++- gigl/distributed/dist_ppr_sampler.py | 29 +- gigl/distributed/utils/dist_typed_sampler.py | 99 ++++-- .../distributed/dist_typed_sampler_test.py | 23 +- 8 files changed, 281 insertions(+), 253 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 25dbc2c8a..1727b1046 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -431,43 +431,13 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, } } -// Helper function for adding one channel's extracted PPR candidates into a -// selection-only candidate list. -// -// This intentionally does not write edge_attr features. Typed extraction uses -// it for the base quota-biased pass, where residual top-up must not affect which -// nodes consume per-channel quota slots. -// -// Inputs: -// nodesAndScores: extracted (node_id, score) candidates for one channel. -// maxScore: largest score in the channel, used to calibrate scores to [0, 1]. -// channelSelectionCandidates: mutable sortable candidates for this channel quota. -// -// Expected output: channelSelectionCandidates contains this channel's calibrated -// candidates for later quota selection. -static void addTypedPPRChannelCandidates(const std::vector>& nodesAndScores, - double maxScore, - std::vector>& channelSelectionCandidates) { - if (nodesAndScores.empty()) { - return; - } - TORCH_CHECK(maxScore > 0.0, - "Typed PPR channel has candidates but non-positive max score ", - maxScore, - ", which indicates invalid PPR state."); - for (const auto& [nodeId, score] : nodesAndScores) { - double calibratedScore = score / maxScore; - channelSelectionCandidates.emplace_back(nodeId, calibratedScore); - } -} - // Helper function for adding one channel's extracted PPR candidates into the -// emitted typed feature table and a fill-pass candidate list. +// emitted typed feature table and a selection candidate list. // -// Unlike addTypedPPRChannelCandidates, this helper owns the output view: it -// populates the edge_attr row for every node it sees. When residual top-up is -// enabled, callers pass residual-aware scores here so base-selected and top-up -// rows are emitted on the same ppr_score + residual scale. +// This helper owns the output view: it populates the edge_attr row 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. @@ -475,7 +445,7 @@ static void addTypedPPRChannelCandidates(const std::vector selectTypedPPRNodeIds( std::vector>>& candidatesByChannel, - const std::vector& channelQuotas, + const std::vector& channelTargetCounts, int32_t maxPPRNodes) { - struct GlobalCandidate { + struct AttributedCandidate { double bestCalibratedScore; double channelCalibratedScore; int32_t channelIndex; int32_t nodeId; }; - // Per-channel ordering is by the channel-local score. Global ordering is by - // the node's best score across all channels, then by the candidate's own - // channel score and stable IDs for deterministic ties. - const auto higherScorePair = [](const auto& a, const auto& b) { - if (a.second != b.second) { - return a.second > b.second; - } - return a.first < b.first; - }; - const auto higherGlobalCandidate = [](const auto& a, const auto& b) { + const auto higherAttributedCandidate = [](const auto& a, const auto& b) { if (a.bestCalibratedScore != b.bestCalibratedScore) { return a.bestCalibratedScore > b.bestCalibratedScore; } @@ -568,93 +532,112 @@ static std::vector selectTypedPPRNodeIds( return a.nodeId < b.nodeId; }; - // Bound reserve sizes by the number of rows that can survive per-channel - // quotas. Current callers already cap channels to these limits, but this - // keeps the helper efficient if a future caller passes larger vectors. - size_t candidateRowReserveSize = 0; - for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByChannel.size()); ++channelIndex) { - candidateRowReserveSize += static_cast( - std::min(channelQuotas[channelIndex], static_cast(candidatesByChannel[channelIndex].size()))); + size_t totalCandidateRows = 0; + for (const auto& candidates : candidatesByChannel) { + totalCandidateRows += candidates.size(); } - // First apply per-channel quotas. candidateRows may still contain the same - // node multiple times if multiple channels reached it, while - // bestCalibratedScoreByNodeId tracks the cross-channel score used for final - // ranking. - std::unordered_map bestCalibratedScoreByNodeId; - bestCalibratedScoreByNodeId.reserve(candidateRowReserveSize); - std::vector candidateRows; - candidateRows.reserve(candidateRowReserveSize); + // 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) { - auto& candidates = candidatesByChannel[channelIndex]; - int32_t channelQuota = channelQuotas[channelIndex]; - int32_t numCandidates = std::min(channelQuota, static_cast(candidates.size())); - if (numCandidates <= 0) { + 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; } - // Only partition when a channel has more rows than its quota. We do not - // need a sorted per-channel prefix; we only need the top quota rows - // before the global cross-channel rank. - if (numCandidates < static_cast(candidates.size())) { - std::nth_element(candidates.begin(), candidates.begin() + numCandidates, candidates.end(), higherScorePair); + 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 < numCandidates; ++candidateIndex) { - int32_t nodeId = candidates[candidateIndex].first; - double calibratedScore = candidates[candidateIndex].second; - candidateRows.push_back({0.0, calibratedScore, channelIndex, nodeId}); - auto [bestScoreIter, inserted] = bestCalibratedScoreByNodeId.emplace(nodeId, calibratedScore); - if (!inserted) { - bestScoreIter->second = std::max(bestScoreIter->second, calibratedScore); - } + for (int32_t candidateIndex = 0; candidateIndex < numTargetCandidates; ++candidateIndex) { + const auto& candidate = candidates[candidateIndex]; + selectedCandidates.push_back(candidate); + selectedNodeIds.insert(candidate.nodeId); } } - // Collapse duplicate node IDs before global ranking. If the same node is - // present through multiple channels, keep the row that would sort highest - // under the final comparator; its bestCalibratedScore is still the best - // score observed for that node across every quota-surviving channel. - std::unordered_map bestCandidateByNodeId; - bestCandidateByNodeId.reserve(bestCalibratedScoreByNodeId.size()); - for (auto candidate : candidateRows) { - candidate.bestCalibratedScore = bestCalibratedScoreByNodeId.at(candidate.nodeId); - auto [candidateIter, inserted] = bestCandidateByNodeId.emplace(candidate.nodeId, candidate); - if (!inserted && higherGlobalCandidate(candidate, candidateIter->second)) { - candidateIter->second = candidate; + // 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); + } } - } - // Move unique candidates into a vector so the final top-k ranking can use - // partial_sort. This avoids sorting the full candidate set when - // maxPPRNodes is smaller than the number of unique channel candidates. - std::vector globalCandidates; - globalCandidates.reserve(bestCandidateByNodeId.size()); - for (const auto& candidateEntry : bestCandidateByNodeId) { - globalCandidates.push_back(candidateEntry.second); + 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]); + } } - int32_t selectedReserveSize = std::min(maxPPRNodes, static_cast(globalCandidates.size())); - if (selectedReserveSize < static_cast(globalCandidates.size())) { - // We only need the best maxPPRNodes rows, so partial_sort ranks the - // returned prefix without paying to fully sort candidates that will be - // discarded. - std::partial_sort(globalCandidates.begin(), - globalCandidates.begin() + selectedReserveSize, - globalCandidates.end(), - higherGlobalCandidate); - } else if (globalCandidates.size() > 1) { - // If every unique candidate is returned, there is no discarded suffix; - // use a regular full sort to preserve the ordered output contract. - std::sort(globalCandidates.begin(), globalCandidates.end(), higherGlobalCandidate); + // 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); } - // globalCandidates is already ranked through selectedReserveSize, so emit - // only node IDs in final order. Dedup already happened above. std::vector selectedNodes; - selectedNodes.reserve(static_cast(selectedReserveSize)); - for (int32_t candidateIndex = 0; candidateIndex < selectedReserveSize; ++candidateIndex) { - selectedNodes.push_back(globalCandidates[candidateIndex].nodeId); + selectedNodes.reserve(selectedCandidates.size()); + for (const auto& candidate : selectedCandidates) { + selectedNodes.push_back(candidate.nodeId); } return selectedNodes; } @@ -705,7 +688,7 @@ PPRExtractResult PPRForwardPush::extractTopKWithResidualTopUp(int32_t maxPPRNode } PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, - const std::vector& channelQuotas, + const std::vector& channelTargetCounts, int32_t maxPPRNodes, bool enableResidualTopUp) { // Typed channels are constructed from the same seed batch and graph schema; @@ -722,16 +705,11 @@ PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector(enableResidualTopUp ? maxPPRNodes : channelPPRNodeBudget); - } + // 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; - std::vector topUpChannelQuotas(static_cast(numChannels), maxPPRNodes); for (int32_t nodeTypeId = 0; nodeTypeId < numNodeTypes; ++nodeTypeId) { std::vector flatIds; @@ -741,45 +719,28 @@ PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector> outputScores; outputScores.reserve(outputCandidateReserveSize); - // Keep the two typed views separate: - // - baseCandidatesByChannel is selection-only. It contains raw - // finalized PPR rows and drives the quota-biased first pass. - // - outputScores/outputCandidatesByChannel is the emitted view. - // It includes residual-aware scores when top-up is enabled, so - // base-selected and fill-selected rows share one score scale. - std::vector>> baseCandidatesByChannel( - static_cast(numChannels)); + // 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]; - int32_t channelPPRNodeBudget = std::min(channelQuotas[channelIndex], maxPPRNodes); - auto baseNodesAndScores = selectFinalizedPPRPairs(nodeTypeState, channelPPRNodeBudget); - - // Keep residual top-up out of the quota-biased base pass. The - // output view starts from the same finalized candidates, then - // adds residual mass/candidates for emitted features and the - // later fill pass. - auto outputNodesAndScores = baseNodesAndScores; + auto outputNodesAndScores = selectFinalizedPPRPairs(nodeTypeState, maxPPRNodes); + if (enableResidualTopUp) { addResidualMassToPPRPairs(nodeTypeState, outputNodesAndScores); appendResidualTopUpPairs(nodeTypeState, outputNodesAndScores, maxPPRNodes); } - double baseMaxScore = 0.0; - for (const auto& nodeAndScore : baseNodesAndScores) { - baseMaxScore = std::max(baseMaxScore, nodeAndScore.second); - } - double outputMaxScore = 0.0; for (const auto& nodeAndScore : outputNodesAndScores) { outputMaxScore = std::max(outputMaxScore, nodeAndScore.second); } - baseCandidatesByChannel[channelIndex].reserve(baseNodesAndScores.size()); outputCandidatesByChannel[channelIndex].reserve(outputNodesAndScores.size()); - addTypedPPRChannelCandidates(baseNodesAndScores, baseMaxScore, baseCandidatesByChannel[channelIndex]); addTypedPPRSeedFeaturesAndCandidates(outputNodesAndScores, outputMaxScore, channelIndex, @@ -788,41 +749,8 @@ PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector(selectedNodes.size()); - if (selectedNodeCount < maxPPRNodes) { - std::unordered_set selectedNodeIds(selectedNodes.begin(), selectedNodes.end()); - auto topUpSelectedNodes = - selectTypedPPRNodeIds(outputCandidatesByChannel, topUpChannelQuotas, maxPPRNodes); - for (int32_t nodeId : topUpSelectedNodes) { - if (selectedNodeCount >= maxPPRNodes) { - break; - } - if (selectedNodeIds.find(nodeId) != selectedNodeIds.end()) { - continue; - } - ++selectedNodeCount; - selectedNodeIds.insert(nodeId); - selectedNodes.push_back(nodeId); - } - } - - if (enableResidualTopUp && selectedNodes.size() > 1) { - std::vector> selectedNodesAndScores; - selectedNodesAndScores.reserve(selectedNodes.size()); - for (int32_t nodeId : selectedNodes) { - selectedNodesAndScores.emplace_back(nodeId, outputScores.at(nodeId)[0]); - } - std::stable_sort(selectedNodesAndScores.begin(), - selectedNodesAndScores.end(), - [](const auto& a, const auto& b) { return a.second > b.second; }); - - selectedNodes.clear(); - selectedNodes.reserve(selectedNodesAndScores.size()); - for (const auto& selectedNodeAndScore : selectedNodesAndScores) { - selectedNodes.push_back(selectedNodeAndScore.first); - } - } for (int32_t nodeId : selectedNodes) { flatIds.push_back(static_cast(nodeId)); diff --git a/gigl-core/core/sampling/ppr_forward_push.h b/gigl-core/core/sampling/ppr_forward_push.h index 0c0765e2b..4fc27f598 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/gigl-core/core/sampling/ppr_forward_push.h @@ -44,9 +44,9 @@ struct SeedNodeTypeState { // 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/channelQuotas vectors. Python - // builds those vectors from typed_channel_quotas insertion order, and this - // function appends indices in ascending order, so the ordering is stable. + // 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 @@ -114,7 +114,7 @@ class PPRForwardPush { PPRExtractResult extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp); friend PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vector& states, - const std::vector& channelQuotas, + const std::vector& channelTargetCounts, int32_t maxPPRNodes, bool enableResidualTopUp); @@ -186,17 +186,20 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector& states, - const std::vector& channelQuotas, + const std::vector& channelTargetCounts, int32_t maxPPRNodes, bool enableResidualTopUp); diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index ec1ca2d4a..9e1fd9f44 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -98,7 +98,7 @@ static py::tuple drainTypedPPRChannelQueuesWrapper(const py::sequence& states, } static PPRExtractResult extractTypedTopKWithResidualTopUpWrapper(const py::sequence& states, - const std::vector& channelQuotas, + const std::vector& channelTargetCounts, int32_t maxPPRNodes, bool enableResidualTopUp) { std::vector statePtrs; @@ -114,7 +114,7 @@ static PPRExtractResult extractTypedTopKWithResidualTopUpWrapper(const py::seque PPRExtractResult result; { py::gil_scoped_release release; - result = extractTypedTopKWithResidualTopUp(statePtrs, channelQuotas, maxPPRNodes, enableResidualTopUp); + result = extractTypedTopKWithResidualTopUp(statePtrs, channelTargetCounts, maxPPRNodes, enableResidualTopUp); } return result; } @@ -149,7 +149,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("extract_typed_top_k_with_residual_top_up", &gigl::extractTypedTopKWithResidualTopUpWrapper, py::arg("states"), - py::arg("channel_quotas"), + py::arg("channel_target_counts"), py::arg("max_ppr_nodes"), py::arg("enable_residual_topup")); } diff --git a/gigl-core/src/gigl_core/ppr_forward_push.pyi b/gigl-core/src/gigl_core/ppr_forward_push.pyi index 44bd82c45..0ec80611a 100644 --- a/gigl-core/src/gigl_core/ppr_forward_push.pyi +++ b/gigl-core/src/gigl_core/ppr_forward_push.pyi @@ -31,7 +31,7 @@ def drain_typed_ppr_channel_queues( ) -> 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_quotas: Sequence[int], + 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 b665b9ae1..ceee2d6e6 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -178,7 +178,7 @@ TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpMergesChannelsInCpp) { channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1})); auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel0, &channel1}, - /*channelQuotas=*/{2, 1}, + /*channelTargetCounts=*/{2, 1}, /*maxPPRNodes=*/3, /*enableResidualTopUp=*/true); @@ -209,6 +209,42 @@ TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpMergesChannelsInCpp) { 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}); @@ -218,7 +254,7 @@ TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpEmitsResidualAwareBaseRows channel.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{1}, /*flatNeighborIds=*/{0}, /*counts=*/{1})); auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel}, - /*channelQuotas=*/{2}, + /*channelTargetCounts=*/{2}, /*maxPPRNodes=*/2, /*enableResidualTopUp=*/true); @@ -249,7 +285,7 @@ TEST(PPRForwardPush, ExtractTypedTopKWithResidualTopUpCanDisableTopUp) { channel1.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0}, /*flatNeighborIds=*/{20}, /*counts=*/{1})); auto result = extractTypedTopKWithResidualTopUp(std::vector{&channel0, &channel1}, - /*channelQuotas=*/{2, 1}, + /*channelTargetCounts=*/{2, 1}, /*maxPPRNodes=*/3, /*enableResidualTopUp=*/false); diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index c14eb34c0..f694afd4a 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -76,8 +76,8 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): 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 quotas, deduplicate shared candidates, and emit the final typed - edge-attribute features. + 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 @@ -167,7 +167,7 @@ def __init__( ] self._is_homogeneous = True - self._typed_ppr_channel_quotas: Optional[list[int]] = None + 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. @@ -385,13 +385,13 @@ def _extract_ppr_state_top_k( def _extract_typed_ppr_state_top_k( self, ppr_states, - typed_ppr_channel_quotas: list[int], + typed_ppr_channel_target_counts: list[int], device: torch.device, ) -> HeteroPPRResult: """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_quotas, + typed_ppr_channel_target_counts, self._max_ppr_nodes, self._enable_residual_topup, ) @@ -532,21 +532,22 @@ async def _compute_typed_ppr_scores( self, seed_nodes: torch.Tensor, seed_node_type: NodeType, - typed_ppr_channel_quotas: list[int], + typed_ppr_channel_target_counts: 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. + 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_quotas: Per-channel candidate quotas, aligned + 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: @@ -633,7 +634,7 @@ def build_ppr_states() -> list[PPRForwardPush]: None, self._extract_typed_ppr_state_top_k, ppr_states, - typed_ppr_channel_quotas, + typed_ppr_channel_target_counts, device, ) @@ -738,8 +739,8 @@ 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()) - typed_ppr_channel_quotas = self._typed_ppr_channel_quotas - if typed_ppr_channel_quotas is None: + 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( @@ -755,7 +756,7 @@ async def _sample_from_nodes( 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_quotas, + typed_ppr_channel_target_counts, ) for seed_type in seed_types ] diff --git a/gigl/distributed/utils/dist_typed_sampler.py b/gigl/distributed/utils/dist_typed_sampler.py index c836746a5..1777b8677 100644 --- a/gigl/distributed/utils/dist_typed_sampler.py +++ b/gigl/distributed/utils/dist_typed_sampler.py @@ -1,11 +1,13 @@ """Construction-time typed-PPR option parsing for distributed samplers. -The helpers in this module validate typed-channel edge-type keys and convert +The helpers in this module validate typed-channel edge-type keys, convert public edge-type keys into the compact integer traversal maps consumed by the -C++ forward-push kernel. They run once during ``DistPPRNeighborSampler`` -initialization, not in the per-batch PPR sampling hot loop. +C++ forward-push kernel, and derive integer channel target counts from public +ratios. They run once during ``DistPPRNeighborSampler`` initialization, not in +the per-batch PPR sampling hot loop. """ +import math from collections.abc import Sequence from typing import Optional, Union, cast @@ -26,7 +28,7 @@ dict[NodeType, torch.Tensor], dict[NodeType, torch.Tensor], ] -# Public typed_channel_quotas keys can be a single edge type or a grouped +# Public typed_channel_ratios keys can be a single edge type or a grouped # channel containing multiple edge types. TypedPPRChannelKey = Union[EdgeType, tuple[EdgeType, ...]] @@ -36,10 +38,10 @@ A tuple of canonical edge types creates one channel whose forward-push state may traverse any edge type in the group. When typed PPR emits multi-column ``edge_attr`` tensors, channel columns follow the insertion order of the -``typed_channel_quotas`` mapping. +``typed_channel_ratios`` mapping. """ # Parsed typed-channel edge-type allowlists, ordered to match the insertion -# order of typed_channel_quotas. +# order of typed_channel_ratios. TypedPPRChannelEdgeTypeGroups = list[tuple[EdgeType, ...]] # One channel's traversal map. The outer list is indexed by integer node-type # ID; each inner list contains the integer edge-type IDs that channel may @@ -49,37 +51,38 @@ TypedPPRChannelTraversalMaps = list[TypedPPRChannelTraversalMap] -def parse_typed_channel_quota_groups( - typed_channel_quotas: Optional[dict[TypedPPRChannelKey, int]], -) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[int]]]: - """Parse typed-PPR channel keys and split keys from quotas. +def parse_typed_channel_ratio_groups( + typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]], +) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[float]]]: + """Parse typed-PPR channel keys and split keys from ratios. Public options allow each channel key to be either one canonical edge type or a non-empty tuple of canonical edge types. Internally, traversal setup needs only the edge-type groups while merge selection needs the aligned - quota values, so this helper returns those two parallel lists. + ratio values, so this helper returns those two parallel lists. This is construction-time option parsing and is not part of the per-batch PPR sampling hot loop. Args: - typed_channel_quotas: User-provided channel mapping from edge-type - allowlist to candidate quota. + typed_channel_ratios: User-provided channel mapping from edge-type + allowlist to target output ratio. Returns: ``(None, None)`` when typed PPR is disabled. Otherwise returns - ``(typed_channel_groups, typed_channel_quota_list)``, both ordered by + ``(typed_channel_groups, typed_channel_ratio_list)``, both ordered by the input mapping insertion order. Raises: ValueError: If a channel key is not a canonical edge type or non-empty - tuple of canonical edge types. + tuple of canonical edge types, or if ratios are not positive and + summing to 1.0. """ - if not typed_channel_quotas: + if not typed_channel_ratios: return None, None typed_channel_groups: TypedPPRChannelEdgeTypeGroups = [] - typed_channel_quota_list: list[int] = [] + typed_channel_ratio_list: list[float] = [] def is_canonical_edge_type(value: object) -> bool: """Return whether ``value`` has PyG's canonical edge-type shape.""" @@ -89,7 +92,7 @@ def is_canonical_edge_type(value: object) -> bool: and all(isinstance(part, str) for part in value) ) - for edge_type_key, quota in typed_channel_quotas.items(): + for edge_type_key, ratio in typed_channel_ratios.items(): if is_canonical_edge_type(edge_type_key): edge_types = (cast(EdgeType, edge_type_key),) elif ( @@ -100,14 +103,64 @@ def is_canonical_edge_type(value: object) -> bool: edge_types = cast(tuple[EdgeType, ...], edge_type_key) else: raise ValueError( - "typed_channel_quotas keys must be a canonical edge type " + "typed_channel_ratios keys must be a canonical edge type " "(src_type, relation, dst_type) or a non-empty tuple of " f"canonical edge types, got {edge_type_key!r}." ) + if isinstance(ratio, bool) or ratio <= 0.0 or ratio > 1.0: + raise ValueError( + "typed_channel_ratios values must be positive ratios in (0, 1], " + f"got {ratio!r} for channel {edge_type_key!r}." + ) typed_channel_groups.append(edge_types) - typed_channel_quota_list.append(quota) + typed_channel_ratio_list.append(float(ratio)) + + ratio_sum = sum(typed_channel_ratio_list) + if not math.isclose(ratio_sum, 1.0, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError( + "typed_channel_ratios values must sum to 1.0, " + f"got {ratio_sum} from ratios {typed_channel_ratio_list}." + ) + + return typed_channel_groups, typed_channel_ratio_list + + +def compute_typed_channel_target_counts( + typed_channel_ratios: list[float], + max_ppr_nodes: int, +) -> list[int]: + """Convert typed-channel ratios to integer per-channel target counts. + + Ratios describe the desired attribution mix in the returned PPR sequence. + This helper converts them to integer counts whose sum is ``max_ppr_nodes``. + Fractional remainders are assigned from largest to smallest, with channel + order as the deterministic tie-breaker. - return typed_channel_groups, typed_channel_quota_list + This is construction-time option parsing and is not part of the per-batch + PPR sampling hot loop. + + Args: + typed_channel_ratios: Per-channel ratios, ordered by typed-channel + insertion order and summing to 1.0. + max_ppr_nodes: Maximum PPR sequence length per seed. + + Returns: + Integer target counts aligned with ``typed_channel_ratios``. + """ + raw_target_counts = [ratio * max_ppr_nodes for ratio in typed_channel_ratios] + target_counts = [math.floor(raw_count) for raw_count in raw_target_counts] + remaining_count = max_ppr_nodes - sum(target_counts) + channels_by_fractional_remainder = sorted( + range(len(raw_target_counts)), + key=lambda channel_index: ( + raw_target_counts[channel_index] - target_counts[channel_index], + -channel_index, + ), + reverse=True, + ) + for channel_index in channels_by_fractional_remainder[:remaining_count]: + target_counts[channel_index] += 1 + return target_counts def build_edge_type_channel_group_edge_type_ids( @@ -150,7 +203,7 @@ def build_edge_type_channel_group_edge_type_ids( unknown_edge_types = set(channel_edge_types) - known_edge_types if unknown_edge_types: raise ValueError( - "typed_channel_quotas includes non-traversable edge types " + "typed_channel_ratios includes non-traversable edge types " f"{sorted(unknown_edge_types)!r}. Known traversable edge " f"types are {sorted(known_edge_types)!r}." ) @@ -169,7 +222,7 @@ def build_edge_type_channel_group_edge_type_ids( ) if not any(node_type_id_to_channel_edge_type_ids): raise ValueError( - "typed_channel_quotas includes edge-type " + "typed_channel_ratios includes edge-type " f"channel={channel_edge_types!r}, " "but no traversable edge types exist for that channel." ) diff --git a/tests/unit/distributed/dist_typed_sampler_test.py b/tests/unit/distributed/dist_typed_sampler_test.py index 0d52f268a..2cc765e92 100644 --- a/tests/unit/distributed/dist_typed_sampler_test.py +++ b/tests/unit/distributed/dist_typed_sampler_test.py @@ -4,7 +4,8 @@ from gigl.distributed.utils.dist_typed_sampler import ( build_edge_type_channel_group_edge_type_ids, - parse_typed_channel_quota_groups, + compute_typed_channel_target_counts, + parse_typed_channel_ratio_groups, ) from tests.test_assets.distributed.test_dataset import ( STORY, @@ -30,16 +31,16 @@ def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( STORY_TO_USER: 1, } - typed_channel_groups, typed_channel_quota_list = ( - parse_typed_channel_quota_groups( + typed_channel_groups, typed_channel_ratio_list = ( + parse_typed_channel_ratio_groups( { - USER_TO_STORY: 2, - (USER_TO_STORY, STORY_TO_USER): 3, + USER_TO_STORY: 0.6, + (USER_TO_STORY, STORY_TO_USER): 0.4, } ) ) assert typed_channel_groups is not None - assert typed_channel_quota_list is not None + assert typed_channel_ratio_list is not None self.assertEqual( typed_channel_groups, @@ -48,7 +49,11 @@ def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( (USER_TO_STORY, STORY_TO_USER), ], ) - self.assertEqual(typed_channel_quota_list, [2, 3]) + self.assertEqual(typed_channel_ratio_list, [0.6, 0.4]) + self.assertEqual( + compute_typed_channel_target_counts(typed_channel_ratio_list, 7), + [4, 3], + ) self.assertEqual( build_edge_type_channel_group_edge_type_ids( edge_type_groups=typed_channel_groups, @@ -63,7 +68,9 @@ def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( ) with self.assertRaisesRegex(ValueError, "canonical edge type"): - parse_typed_channel_quota_groups({("bad",): 1}) + parse_typed_channel_ratio_groups({("bad",): 1.0}) + with self.assertRaisesRegex(ValueError, "sum to 1.0"): + parse_typed_channel_ratio_groups({USER_TO_STORY: 0.5}) with self.assertRaisesRegex(ValueError, "non-traversable edge types"): build_edge_type_channel_group_edge_type_ids( edge_type_groups=[(("unknown", "edge", "type"),)], From 13b409195a598f0c1e93934f07c3a06cb2d2ae6a Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 30 Jul 2026 07:11:51 +0000 Subject: [PATCH 06/10] Move typed PPR parsing helpers out of C++ PR --- gigl/distributed/dist_ppr_sampler.py | 24 +- gigl/distributed/utils/dist_typed_sampler.py | 230 ------------------ .../distributed/dist_typed_sampler_test.py | 84 ------- 3 files changed, 19 insertions(+), 319 deletions(-) delete mode 100644 gigl/distributed/utils/dist_typed_sampler.py delete mode 100644 tests/unit/distributed/dist_typed_sampler_test.py diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index f694afd4a..02e566107 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -21,11 +21,6 @@ from graphlearn_torch.utils import merge_dict from gigl.distributed.base_sampler import BaseDistNeighborSampler -from gigl.distributed.utils.dist_typed_sampler import ( - HeteroPPRResult, - PPRResult, - TypedPPRChannelTraversalMaps, -) from gigl.types.graph import DEFAULT_HOMOGENEOUS_NODE_TYPE, is_label_edge_type # Trailing "." is an intentional separator. These constants are used both to @@ -45,6 +40,25 @@ DEFAULT_HOMOGENEOUS_NODE_TYPE, ) +# C++ PPR extraction output: flat node IDs, flat weights, and per-seed valid +# counts. Homogeneous extraction uses tensors directly; heterogeneous extraction +# uses dictionaries keyed by node type. +PPRResult = tuple[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], +] +# Heterogeneous-only view of PPRResult after typed PPR extraction. +HeteroPPRResult = tuple[ + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], + dict[NodeType, torch.Tensor], +] +# All typed-channel traversal maps. Each channel map is indexed first by integer +# node-type ID, then contains the integer edge-type IDs that channel may traverse +# from that node type. +TypedPPRChannelTraversalMaps = list[list[list[int]]] + class DistPPRNeighborSampler(BaseDistNeighborSampler): """Personalized PageRank (PPR) based distributed neighbor sampler. diff --git a/gigl/distributed/utils/dist_typed_sampler.py b/gigl/distributed/utils/dist_typed_sampler.py deleted file mode 100644 index 1777b8677..000000000 --- a/gigl/distributed/utils/dist_typed_sampler.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Construction-time typed-PPR option parsing for distributed samplers. - -The helpers in this module validate typed-channel edge-type keys, convert -public edge-type keys into the compact integer traversal maps consumed by the -C++ forward-push kernel, and derive integer channel target counts from public -ratios. They run once during ``DistPPRNeighborSampler`` initialization, not in -the per-batch PPR sampling hot loop. -""" - -import math -from collections.abc import Sequence -from typing import Optional, Union, cast - -import torch -from graphlearn_torch.typing import EdgeType, NodeType - -# C++ PPR extraction output: flat node IDs, flat weights, and per-seed valid -# counts. Homogeneous extraction uses tensors directly; heterogeneous extraction -# uses dictionaries keyed by node type. -PPRResult = tuple[ - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], -] -# Heterogeneous-only view of PPRResult after typed PPR extraction. -HeteroPPRResult = tuple[ - dict[NodeType, torch.Tensor], - dict[NodeType, torch.Tensor], - dict[NodeType, torch.Tensor], -] -# Public typed_channel_ratios keys can be a single edge type or a grouped -# channel containing multiple edge types. -TypedPPRChannelKey = Union[EdgeType, tuple[EdgeType, ...]] - -"""TypedPPRChannelKey describes one public typed-PPR traversal channel key. - -A single canonical edge type creates one channel restricted to that edge type. -A tuple of canonical edge types creates one channel whose forward-push state may -traverse any edge type in the group. When typed PPR emits multi-column -``edge_attr`` tensors, channel columns follow the insertion order of the -``typed_channel_ratios`` mapping. -""" -# Parsed typed-channel edge-type allowlists, ordered to match the insertion -# order of typed_channel_ratios. -TypedPPRChannelEdgeTypeGroups = list[tuple[EdgeType, ...]] -# One channel's traversal map. The outer list is indexed by integer node-type -# ID; each inner list contains the integer edge-type IDs that channel may -# traverse from that node type. -TypedPPRChannelTraversalMap = list[list[int]] -# All typed-channel traversal maps, ordered to match typed-channel order. -TypedPPRChannelTraversalMaps = list[TypedPPRChannelTraversalMap] - - -def parse_typed_channel_ratio_groups( - typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]], -) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[float]]]: - """Parse typed-PPR channel keys and split keys from ratios. - - Public options allow each channel key to be either one canonical edge type - or a non-empty tuple of canonical edge types. Internally, traversal setup - needs only the edge-type groups while merge selection needs the aligned - ratio values, so this helper returns those two parallel lists. - - This is construction-time option parsing and is not part of the per-batch - PPR sampling hot loop. - - Args: - typed_channel_ratios: User-provided channel mapping from edge-type - allowlist to target output ratio. - - Returns: - ``(None, None)`` when typed PPR is disabled. Otherwise returns - ``(typed_channel_groups, typed_channel_ratio_list)``, both ordered by - the input mapping insertion order. - - Raises: - ValueError: If a channel key is not a canonical edge type or non-empty - tuple of canonical edge types, or if ratios are not positive and - summing to 1.0. - """ - if not typed_channel_ratios: - return None, None - - typed_channel_groups: TypedPPRChannelEdgeTypeGroups = [] - typed_channel_ratio_list: list[float] = [] - - def is_canonical_edge_type(value: object) -> bool: - """Return whether ``value`` has PyG's canonical edge-type shape.""" - return ( - isinstance(value, tuple) - and len(value) == 3 - and all(isinstance(part, str) for part in value) - ) - - for edge_type_key, ratio in typed_channel_ratios.items(): - if is_canonical_edge_type(edge_type_key): - edge_types = (cast(EdgeType, edge_type_key),) - elif ( - isinstance(edge_type_key, tuple) - and edge_type_key - and all(is_canonical_edge_type(edge_type) for edge_type in edge_type_key) - ): - edge_types = cast(tuple[EdgeType, ...], edge_type_key) - else: - raise ValueError( - "typed_channel_ratios keys must be a canonical edge type " - "(src_type, relation, dst_type) or a non-empty tuple of " - f"canonical edge types, got {edge_type_key!r}." - ) - if isinstance(ratio, bool) or ratio <= 0.0 or ratio > 1.0: - raise ValueError( - "typed_channel_ratios values must be positive ratios in (0, 1], " - f"got {ratio!r} for channel {edge_type_key!r}." - ) - typed_channel_groups.append(edge_types) - typed_channel_ratio_list.append(float(ratio)) - - ratio_sum = sum(typed_channel_ratio_list) - if not math.isclose(ratio_sum, 1.0, rel_tol=1e-9, abs_tol=1e-9): - raise ValueError( - "typed_channel_ratios values must sum to 1.0, " - f"got {ratio_sum} from ratios {typed_channel_ratio_list}." - ) - - return typed_channel_groups, typed_channel_ratio_list - - -def compute_typed_channel_target_counts( - typed_channel_ratios: list[float], - max_ppr_nodes: int, -) -> list[int]: - """Convert typed-channel ratios to integer per-channel target counts. - - Ratios describe the desired attribution mix in the returned PPR sequence. - This helper converts them to integer counts whose sum is ``max_ppr_nodes``. - Fractional remainders are assigned from largest to smallest, with channel - order as the deterministic tie-breaker. - - This is construction-time option parsing and is not part of the per-batch - PPR sampling hot loop. - - Args: - typed_channel_ratios: Per-channel ratios, ordered by typed-channel - insertion order and summing to 1.0. - max_ppr_nodes: Maximum PPR sequence length per seed. - - Returns: - Integer target counts aligned with ``typed_channel_ratios``. - """ - raw_target_counts = [ratio * max_ppr_nodes for ratio in typed_channel_ratios] - target_counts = [math.floor(raw_count) for raw_count in raw_target_counts] - remaining_count = max_ppr_nodes - sum(target_counts) - channels_by_fractional_remainder = sorted( - range(len(raw_target_counts)), - key=lambda channel_index: ( - raw_target_counts[channel_index] - target_counts[channel_index], - -channel_index, - ), - reverse=True, - ) - for channel_index in channels_by_fractional_remainder[:remaining_count]: - target_counts[channel_index] += 1 - return target_counts - - -def build_edge_type_channel_group_edge_type_ids( - edge_type_groups: TypedPPRChannelEdgeTypeGroups, - edge_type_to_edge_type_id: dict[EdgeType, int], - node_type_to_edge_types: dict[NodeType, list[EdgeType]], - node_types: Sequence[NodeType], -) -> TypedPPRChannelTraversalMaps: - """Convert typed-channel edge-type allowlists to PPRForwardPush IDs. - - Returns one traversal map per typed channel, ordered to match - ``edge_type_groups``. A single traversal map has shape - ``list[list[int]]``: the outer index is ``node_type_id``, and the inner list - contains allowed ``edge_type_id`` values for that node type in that channel. - - This conversion runs once during sampler construction; the resulting integer - maps are reused by the per-batch C++ PPR states. - - Args: - edge_type_groups: Ordered typed channels, where each channel is the - canonical edge types that its PPR state may traverse. - edge_type_to_edge_type_id: Mapping from canonical edge type to the - compact integer ID used by the C++ forward-push kernel. - node_type_to_edge_types: Traversable edge types keyed by anchor node - type, after label-edge filtering and edge-direction handling. - node_types: Ordered node types whose positions match the kernel's - integer node-type IDs. - - Returns: - ``channel_traversal_maps[channel_id][node_type_id]`` gives the allowed - integer edge-type IDs that channel may traverse from that node type. - - Raises: - ValueError: If a configured edge type is unknown, excluded from PPR - traversal, or cannot be traversed from any node type. - """ - known_edge_types = set(edge_type_to_edge_type_id.keys()) - channel_edge_type_ids_by_node_type: TypedPPRChannelTraversalMaps = [] - for channel_edge_types in edge_type_groups: - unknown_edge_types = set(channel_edge_types) - known_edge_types - if unknown_edge_types: - raise ValueError( - "typed_channel_ratios includes non-traversable edge types " - f"{sorted(unknown_edge_types)!r}. Known traversable edge " - f"types are {sorted(known_edge_types)!r}." - ) - - channel_edge_type_set = set(channel_edge_types) - node_type_id_to_channel_edge_type_ids: TypedPPRChannelTraversalMap = [] - for node_type in node_types: - channel_edge_type_ids_for_node_type: list[int] = [] - for edge_type in node_type_to_edge_types.get(node_type, []): - if edge_type in channel_edge_type_set: - channel_edge_type_ids_for_node_type.append( - edge_type_to_edge_type_id[edge_type] - ) - node_type_id_to_channel_edge_type_ids.append( - channel_edge_type_ids_for_node_type - ) - if not any(node_type_id_to_channel_edge_type_ids): - raise ValueError( - "typed_channel_ratios includes edge-type " - f"channel={channel_edge_types!r}, " - "but no traversable edge types exist for that channel." - ) - channel_edge_type_ids_by_node_type.append(node_type_id_to_channel_edge_type_ids) - return channel_edge_type_ids_by_node_type diff --git a/tests/unit/distributed/dist_typed_sampler_test.py b/tests/unit/distributed/dist_typed_sampler_test.py deleted file mode 100644 index 2cc765e92..000000000 --- a/tests/unit/distributed/dist_typed_sampler_test.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Unit tests for typed-PPR sampler construction helpers.""" - -from absl.testing import absltest - -from gigl.distributed.utils.dist_typed_sampler import ( - build_edge_type_channel_group_edge_type_ids, - compute_typed_channel_target_counts, - parse_typed_channel_ratio_groups, -) -from tests.test_assets.distributed.test_dataset import ( - STORY, - STORY_TO_USER, - USER, - USER_TO_STORY, -) -from tests.test_assets.test_case import TestCase - - -class DistTypedSamplerTest(TestCase): - def test_typed_ppr_edge_type_channels_parse_and_build_traversal_maps( - self, - ) -> None: - """Verify typed-PPR can use canonical edge-type channels.""" - node_type_to_edge_types = { - USER: [USER_TO_STORY], - STORY: [STORY_TO_USER], - } - node_types = [USER, STORY] - edge_type_to_edge_type_id = { - USER_TO_STORY: 0, - STORY_TO_USER: 1, - } - - typed_channel_groups, typed_channel_ratio_list = ( - parse_typed_channel_ratio_groups( - { - USER_TO_STORY: 0.6, - (USER_TO_STORY, STORY_TO_USER): 0.4, - } - ) - ) - assert typed_channel_groups is not None - assert typed_channel_ratio_list is not None - - self.assertEqual( - typed_channel_groups, - [ - (USER_TO_STORY,), - (USER_TO_STORY, STORY_TO_USER), - ], - ) - self.assertEqual(typed_channel_ratio_list, [0.6, 0.4]) - self.assertEqual( - compute_typed_channel_target_counts(typed_channel_ratio_list, 7), - [4, 3], - ) - self.assertEqual( - build_edge_type_channel_group_edge_type_ids( - edge_type_groups=typed_channel_groups, - edge_type_to_edge_type_id=edge_type_to_edge_type_id, - node_type_to_edge_types=node_type_to_edge_types, - node_types=node_types, - ), - [ - [[0], []], - [[0], [1]], - ], - ) - - with self.assertRaisesRegex(ValueError, "canonical edge type"): - parse_typed_channel_ratio_groups({("bad",): 1.0}) - with self.assertRaisesRegex(ValueError, "sum to 1.0"): - parse_typed_channel_ratio_groups({USER_TO_STORY: 0.5}) - with self.assertRaisesRegex(ValueError, "non-traversable edge types"): - build_edge_type_channel_group_edge_type_ids( - edge_type_groups=[(("unknown", "edge", "type"),)], - edge_type_to_edge_type_id=edge_type_to_edge_type_id, - node_type_to_edge_types=node_type_to_edge_types, - node_types=node_types, - ) - - -if __name__ == "__main__": - absltest.main() From 009065f349695a4fbdde2a4934c2140c377a7849 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 30 Jul 2026 07:28:06 +0000 Subject: [PATCH 07/10] Address typed PPR C++ review comments --- gigl-core/core/sampling/ppr_forward_push.cpp | 29 ++++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 1727b1046..974eeb5e0 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -434,10 +434,10 @@ static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState, // Helper function for adding one channel's extracted PPR candidates into the // emitted typed feature table and a selection candidate list. // -// This helper owns the output view: it populates the edge_attr row 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. +// 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. @@ -485,8 +485,9 @@ static void addTypedPPRSeedFeaturesAndCandidates(const std::vector_batchSize; int32_t numNodeTypes = firstState->_numNodeTypes; @@ -735,10 +738,12 @@ PPRExtractResult extractTypedTopKWithResidualTopUp(const std::vectorsecond : 0.0; outputCandidatesByChannel[channelIndex].reserve(outputNodesAndScores.size()); addTypedPPRSeedFeaturesAndCandidates(outputNodesAndScores, From 4d799f025b203b55480b05dc63a0de058fd38235 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 30 Jul 2026 07:42:57 +0000 Subject: [PATCH 08/10] Simplify typed PPR channel score write --- gigl-core/core/sampling/ppr_forward_push.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 974eeb5e0..990d65332 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -486,9 +486,8 @@ static void addTypedPPRSeedFeaturesAndCandidates(const std::vector Date: Thu, 30 Jul 2026 07:58:36 +0000 Subject: [PATCH 09/10] Parallelize typed PPR channel drains --- gigl-core/core/sampling/ppr_forward_push.cpp | 27 +++++++++++++++----- gigl-core/core/sampling/ppr_forward_push.h | 10 +++++--- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 990d65332..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 @@ -163,14 +164,26 @@ TypedPPRQueueDrainResult drainTypedPPRChannelQueues(const std::vector> unionedSourceNodeIdsByEdgeTypeId; - // TODO: If benchmarking shows typed PPR is materially slower than regular - // heterogeneous PPR because of this drain loop, evaluate parallelizing - // channels with per-thread frontier maps and a deterministic merge into - // queueDrainResult. - for (size_t channelIndex = 0; channelIndex < states.size(); ++channelIndex) { - PPRForwardPush* state = states[channelIndex]; + 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(); + } + } - auto channelFrontierByEdgeTypeId = state->drainQueue(); + // 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; } diff --git a/gigl-core/core/sampling/ppr_forward_push.h b/gigl-core/core/sampling/ppr_forward_push.h index 4fc27f598..e909fa479 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/gigl-core/core/sampling/ppr_forward_push.h @@ -162,10 +162,12 @@ class PPRForwardPush { // Helper function for draining several independent channel states for one // typed-PPR iteration. // -// This is the typed wrapper around PPRForwardPush::drainQueue(): it drains each -// channel state, 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. +// 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 From 006f24e65beae04b7aed4328794d9944c78911b7 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 30 Jul 2026 08:15:50 +0000 Subject: [PATCH 10/10] Inline lower PR PPR result annotations --- gigl/distributed/dist_ppr_sampler.py | 47 ++++++++++++++-------------- 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 02e566107..bcffb5ff0 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -40,25 +40,6 @@ DEFAULT_HOMOGENEOUS_NODE_TYPE, ) -# C++ PPR extraction output: flat node IDs, flat weights, and per-seed valid -# counts. Homogeneous extraction uses tensors directly; heterogeneous extraction -# uses dictionaries keyed by node type. -PPRResult = tuple[ - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], - Union[torch.Tensor, dict[NodeType, torch.Tensor]], -] -# Heterogeneous-only view of PPRResult after typed PPR extraction. -HeteroPPRResult = tuple[ - dict[NodeType, torch.Tensor], - dict[NodeType, torch.Tensor], - dict[NodeType, torch.Tensor], -] -# All typed-channel traversal maps. Each channel map is indexed first by integer -# node-type ID, then contains the integer edge-type IDs that channel may traverse -# from that node type. -TypedPPRChannelTraversalMaps = list[list[list[int]]] - class DistPPRNeighborSampler(BaseDistNeighborSampler): """Personalized PageRank (PPR) based distributed neighbor sampler. @@ -245,7 +226,9 @@ def __init__( for node_type in all_node_types ] - self._typed_ppr_channel_to_node_type_id_to_edge_type_ids: TypedPPRChannelTraversalMaps = [] + self._typed_ppr_channel_to_node_type_id_to_edge_type_ids: list[ + list[list[int]] + ] = [] def _convert_degree_tensors_to_dict( self, @@ -342,7 +325,11 @@ def _extract_ppr_state_top_k( self, ppr_state, device: torch.device, - ) -> PPRResult: + ) -> tuple[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + ]: """Extract PPR neighbors from a completed C++ Forward Push state. The C++ kernel indexes node types by compact integer IDs for speed. @@ -401,7 +388,11 @@ def _extract_typed_ppr_state_top_k( ppr_states, typed_ppr_channel_target_counts: list[int], device: torch.device, - ) -> HeteroPPRResult: + ) -> 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, @@ -431,7 +422,11 @@ async def _compute_ppr_scores( self, seed_nodes: torch.Tensor, seed_node_type: Optional[NodeType] = None, - ) -> PPRResult: + ) -> tuple[ + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + Union[torch.Tensor, dict[NodeType, torch.Tensor]], + ]: """ Compute PPR scores for seed nodes using the push-based approximation algorithm. @@ -547,7 +542,11 @@ async def _compute_typed_ppr_scores( seed_nodes: torch.Tensor, seed_node_type: NodeType, typed_ppr_channel_target_counts: list[int], - ) -> HeteroPPRResult: + ) -> 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