diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index a2ab255a8..770a37045 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -4,7 +4,9 @@ #include #include +#include #include +#include #include #include #include @@ -147,6 +149,65 @@ std::optional> PPRForwardPush::drainQ return result; } +DrainedTypedPPRQueues 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."); + + DrainedTypedPPRQueues drained; + std::unordered_map> nodeIdsByEdgeTypeId; + + for (size_t channelIndex = 0; channelIndex < states.size(); ++channelIndex) { + PPRForwardPush* state = states[channelIndex]; + TORCH_CHECK(state != nullptr, "PPR state pointer must not be null."); + + auto nodesByEdgeTypeId = state->drainQueue(); + if (!nodesByEdgeTypeId.has_value()) { + continue; + } + + drained.activeChannelIndices.push_back(static_cast(channelIndex)); + + bool fetchBudgetRemaining = maxFetchIterations < 0 || fetchIterationCounts[channelIndex] < maxFetchIterations; + if (!fetchBudgetRemaining) { + continue; + } + + std::vector requestedEdgeTypeIds; + for (const auto& [edgeTypeId, nodes] : nodesByEdgeTypeId.value()) { + if (nodes.numel() == 0) { + continue; + } + TORCH_CHECK(nodes.dim() == 1, "drainQueue() must return 1-D node tensors."); + TORCH_CHECK(nodes.scalar_type() == torch::kInt64, "drainQueue() must return int64 node tensors."); + + requestedEdgeTypeIds.push_back(edgeTypeId); + auto contiguousNodes = nodes.contiguous(); + auto nodeAccessor = contiguousNodes.accessor(); + auto& nodeIds = nodeIdsByEdgeTypeId[edgeTypeId]; + for (int64_t nodeIndex = 0; nodeIndex < contiguousNodes.size(0); ++nodeIndex) { + nodeIds.insert(nodeAccessor[nodeIndex]); + } + } + + if (!requestedEdgeTypeIds.empty()) { + drained.fetchChannelIndices.push_back(static_cast(channelIndex)); + drained.edgeTypeIdsByFetchChannel.push_back(std::move(requestedEdgeTypeIds)); + } + } + + for (const auto& [edgeTypeId, nodeIds] : nodeIdsByEdgeTypeId) { + std::vector nodeIdsToLookup(nodeIds.begin(), nodeIds.end()); + drained.unionNodesByEdgeTypeId[edgeTypeId] = torch::tensor(nodeIdsToLookup, torch::kLong); + } + return drained; +} + void PPRForwardPush::pushResiduals( const std::unordered_map>& fetchedByEtypeId) { // Step 1: Persist fetched neighbor lists in the per-state cache. drainQueue() @@ -266,9 +327,196 @@ void PPRForwardPush::pushResiduals( } } +static std::vector> selectPPRPairsWithResidualTopUp(const SeedNodeTypeState& nodeTypeState, + int32_t maxPPRNodes, + int32_t maxResidualNodes, + int32_t maxTotalNodes) { + const auto& scores = nodeTypeState.pprScores; + const auto higherScore = [](const auto& a, const auto& b) { return a.second > b.second; }; + + int32_t totalNodeLimit = maxPPRNodes + maxResidualNodes; + if (maxTotalNodes >= 0) { + totalNodeLimit = maxTotalNodes; + } + + const int32_t topK = std::min(maxPPRNodes, static_cast(scores.size())); + const int32_t residualTopUpBudget = std::min(maxResidualNodes, std::max(0, totalNodeLimit - topK)); + std::vector> selectedPairs; + selectedPairs.reserve(static_cast(topK + residualTopUpBudget)); + std::unordered_set selectedPPRNodeIds; + if (topK > 0) { + if (residualTopUpBudget > 0) { + selectedPPRNodeIds.reserve(static_cast(topK)); + } + std::vector> scorePairs(scores.begin(), scores.end()); + // Selection is intentionally two-phase: finalized nodes are selected + // first by raw PPR score, and residual candidates only compete for + // the remaining budget. + if (maxResidualNodes > 0) { + // The final emitted order is sorted by ppr_score + residual after + // top-up candidates are selected, so this pass only needs to + // partition out the raw-PPR top K. + if (topK < static_cast(scorePairs.size())) { + std::nth_element(scorePairs.begin(), scorePairs.begin() + topK, scorePairs.end(), higherScore); + } + } else { + std::partial_sort(scorePairs.begin(), scorePairs.begin() + topK, scorePairs.end(), higherScore); + } + + for (int32_t rankIdx = 0; rankIdx < topK; ++rankIdx) { + int32_t nodeId = scorePairs[rankIdx].first; + double outputScore = scorePairs[rankIdx].second; + if (maxResidualNodes > 0) { + auto residualIter = nodeTypeState.residuals.find(nodeId); + if (residualIter != nodeTypeState.residuals.end()) { + outputScore += residualIter->second; + } + } + selectedPairs.emplace_back(nodeId, outputScore); + if (residualTopUpBudget > 0) { + selectedPPRNodeIds.insert(nodeId); + } + } + } + + if (residualTopUpBudget > 0) { + std::vector> residualPairs; + residualPairs.reserve(nodeTypeState.residuals.size()); + for (const auto& [nodeId, residual] : nodeTypeState.residuals) { + if (residual <= 0.0 || selectedPPRNodeIds.find(nodeId) != selectedPPRNodeIds.end()) { + continue; + } + + auto scoreIter = scores.find(nodeId); + double pprScore = (scoreIter != scores.end()) ? scoreIter->second : 0.0; + double outputScore = pprScore + residual; + residualPairs.emplace_back(nodeId, outputScore); + } + + const int32_t residualTopK = std::min(residualTopUpBudget, static_cast(residualPairs.size())); + if (residualTopK > 0) { + // Residual candidates only need selection here; selected finalized + // and residual rows are sorted together below. + if (residualTopK < static_cast(residualPairs.size())) { + std::nth_element( + residualPairs.begin(), residualPairs.begin() + residualTopK, residualPairs.end(), higherScore); + } + + for (int32_t rankIdx = 0; rankIdx < residualTopK; ++rankIdx) { + selectedPairs.emplace_back(residualPairs[rankIdx].first, residualPairs[rankIdx].second); + } + } + } + + if (maxResidualNodes > 0 && selectedPairs.size() > 1) { + std::sort(selectedPairs.begin(), selectedPairs.end(), higherScore); + } + return selectedPairs; +} + +static double clampTypedPPRScore(double score) { + if (!std::isfinite(score)) { + return std::numeric_limits::quiet_NaN(); + } + return std::min(std::max(score, 0.0), 1.0); +} + +static void addTypedPPRSeedCandidates(std::unordered_map>& seedScores, + std::vector>& seedChannelCandidates, + const std::vector>& nodesAndScores, + double maxScore, + int32_t channelIndex, + int32_t numEdgeAttrFeatures, + int32_t numChannels) { + for (const auto& [nodeId, rawScore] : nodesAndScores) { + double score = clampTypedPPRScore(rawScore); + if (!std::isfinite(score)) { + continue; + } + double calibratedScore = maxScore > 0.0 ? score / maxScore : 0.0; + auto scoreIter = seedScores.find(nodeId); + if (scoreIter == seedScores.end()) { + scoreIter = seedScores.emplace(nodeId, std::vector(numEdgeAttrFeatures, 0.0)).first; + } + auto& scoreFeatures = scoreIter->second; + scoreFeatures[0] = std::max(scoreFeatures[0], calibratedScore); + int32_t channelScoreIndex = 1 + channelIndex; + int32_t channelPresenceIndex = 1 + numChannels + channelIndex; + scoreFeatures[channelScoreIndex] = std::max(scoreFeatures[channelScoreIndex], calibratedScore); + scoreFeatures[channelPresenceIndex] = 1.0; + seedChannelCandidates.emplace_back(nodeId, calibratedScore); + } +} + +static std::vector selectTypedPPRNodeIds( + const std::unordered_map>& seedScores, + std::vector>> candidatesByChannel, + const std::vector& channelQuotas, + int32_t maxPPRNodes) { + struct GlobalCandidate { + double bestCalibratedScore; + double channelCalibratedScore; + int32_t channelIndex; + int32_t nodeId; + }; + + std::vector globalCandidates; + for (int32_t channelIndex = 0; channelIndex < static_cast(candidatesByChannel.size()); ++channelIndex) { + auto& candidates = candidatesByChannel[channelIndex]; + std::sort(candidates.begin(), candidates.end(), [](const auto& a, const auto& b) { + if (a.second != b.second) { + return a.second > b.second; + } + return a.first < b.first; + }); + + int32_t channelQuota = channelQuotas[channelIndex]; + int32_t numCandidates = std::min(channelQuota, static_cast(candidates.size())); + for (int32_t candidateIndex = 0; candidateIndex < numCandidates; ++candidateIndex) { + int32_t nodeId = candidates[candidateIndex].first; + auto scoreIter = seedScores.find(nodeId); + if (scoreIter == seedScores.end()) { + continue; + } + globalCandidates.push_back({scoreIter->second[0], candidates[candidateIndex].second, channelIndex, nodeId}); + } + } + + std::sort(globalCandidates.begin(), globalCandidates.end(), [](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; + }); + + std::unordered_set selectedNodeIds; + std::vector selectedNodes; + int32_t selectedReserveSize = std::min(maxPPRNodes, static_cast(globalCandidates.size())); + selectedNodeIds.reserve(static_cast(selectedReserveSize)); + selectedNodes.reserve(static_cast(selectedReserveSize)); + for (const auto& candidate : globalCandidates) { + if (static_cast(selectedNodes.size()) >= maxPPRNodes) { + break; + } + if (selectedNodeIds.find(candidate.nodeId) != selectedNodeIds.end()) { + continue; + } + selectedNodeIds.insert(candidate.nodeId); + selectedNodes.push_back(candidate.nodeId); + } + return selectedNodes; +} + std::unordered_map> PPRForwardPush:: extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) { TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, "."); + int32_t residualTopUpNodes = enableResidualTopUp ? maxPPRNodes : 0; std::unordered_map> result; // Emit an entry for every node type, even if unreachable in this batch (empty tensors, @@ -281,94 +529,163 @@ std::unordered_map b.second; }; - - const int32_t topK = std::min(maxPPRNodes, static_cast(scores.size())); - const int32_t residualTopUpBudget = enableResidualTopUp ? maxPPRNodes - topK : 0; - std::vector> selectedPairs; - selectedPairs.reserve(static_cast(topK) + static_cast(residualTopUpBudget)); - std::unordered_set selectedPPRNodeIds; - if (topK > 0) { - if (residualTopUpBudget > 0) { - selectedPPRNodeIds.reserve(static_cast(topK)); - } - std::vector> scorePairs(scores.begin(), scores.end()); - // Selection is intentionally two-phase: finalized nodes are selected - // first by raw PPR score, and residual candidates only compete for - // the remaining budget. - if (enableResidualTopUp) { - // The final emitted order is sorted by ppr_score + residual - // after top-up candidates are selected, so this pass only - // needs to partition out the raw-PPR top K. - if (topK < static_cast(scorePairs.size())) { - std::nth_element(scorePairs.begin(), scorePairs.begin() + topK, scorePairs.end(), higherScore); - } - } else { - std::partial_sort(scorePairs.begin(), scorePairs.begin() + topK, scorePairs.end(), higherScore); - } + auto selectedPairs = + selectPPRPairsWithResidualTopUp(nodeTypeState, maxPPRNodes, residualTopUpNodes, maxPPRNodes); + for (const auto& [nodeId, score] : selectedPairs) { + flatIds.push_back(static_cast(nodeId)); + flatWeights.push_back(score); + } + validCounts.push_back(static_cast(selectedPairs.size())); + } - for (int32_t rankIdx = 0; rankIdx < topK; ++rankIdx) { - int32_t nodeId = scorePairs[rankIdx].first; - double outputScore = scorePairs[rankIdx].second; - if (enableResidualTopUp) { - auto residualIter = nodeTypeState.residuals.find(nodeId); - if (residualIter != nodeTypeState.residuals.end()) { - outputScore += residualIter->second; - } + result[nodeTypeId] = {torch::tensor(flatIds, torch::kLong), + torch::tensor(flatWeights, torch::kDouble), + torch::tensor(validCounts, torch::kLong)}; + } + return result; +} + +std::unordered_map> extractTypedTopKWithResidualTopUp( + const std::vector& states, + const std::vector& channelQuotas, + int32_t maxPPRNodes, + bool enableResidualTopUp) { + TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, "."); + TORCH_CHECK(!states.empty(), "Typed PPR extraction requires at least one channel state."); + TORCH_CHECK(states.size() == channelQuotas.size(), + "Expected one channel quota per typed PPR state, got ", + channelQuotas.size(), + " quotas for ", + states.size(), + " states."); + + const auto* firstState = states.front(); + TORCH_CHECK(firstState != nullptr, "PPR state pointer must not be null."); + int32_t batchSize = firstState->_batchSize; + int32_t numNodeTypes = firstState->_numNodeTypes; + int32_t numChannels = static_cast(states.size()); + int32_t numEdgeAttrFeatures = 1 + (2 * numChannels); + int32_t residualTopUpNodes = enableResidualTopUp ? maxPPRNodes : 0; + + for (int32_t channelIndex = 0; channelIndex < numChannels; ++channelIndex) { + TORCH_CHECK(channelQuotas[channelIndex] >= 0, + "channelQuotas must be non-negative, got ", + channelQuotas[channelIndex], + " at channel ", + channelIndex, + "."); + TORCH_CHECK(states[channelIndex] != nullptr, "PPR state pointer must not be null."); + TORCH_CHECK(states[channelIndex]->_batchSize == batchSize, + "All typed PPR channel states must have the same batch size."); + TORCH_CHECK(states[channelIndex]->_numNodeTypes == numNodeTypes, + "All typed PPR channel states must have the same node-type count."); + } + + std::unordered_map> result; + 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> baseScores; + std::unordered_map> extendedScores; + std::vector>> baseCandidatesByChannel( + static_cast(numChannels)); + std::vector>> extendedCandidatesByChannel( + static_cast(numChannels)); + + for (int32_t channelIndex = 0; channelIndex < numChannels; ++channelIndex) { + const auto& nodeTypeState = states[channelIndex]->_state[seedIdx][nodeTypeId]; + int32_t channelPPRLimit = std::min(channelQuotas[channelIndex], maxPPRNodes); + auto selectedPairs = + selectPPRPairsWithResidualTopUp(nodeTypeState, channelPPRLimit, residualTopUpNodes, maxPPRNodes); + + double extendedMaxScore = 0.0; + std::vector> baseNodesAndScores; + std::vector> extendedNodesAndScores; + baseNodesAndScores.reserve(static_cast( + std::min(channelQuotas[channelIndex], static_cast(selectedPairs.size())))); + extendedNodesAndScores.reserve(selectedPairs.size()); + for (int32_t candidateIndex = 0; candidateIndex < static_cast(selectedPairs.size()); + ++candidateIndex) { + double score = clampTypedPPRScore(selectedPairs[candidateIndex].second); + if (!std::isfinite(score)) { + continue; } - selectedPairs.emplace_back(nodeId, outputScore); - if (residualTopUpBudget > 0) { - selectedPPRNodeIds.insert(nodeId); + extendedNodesAndScores.emplace_back(selectedPairs[candidateIndex].first, score); + extendedMaxScore = std::max(extendedMaxScore, score); + if (candidateIndex < channelQuotas[channelIndex]) { + baseNodesAndScores.emplace_back(selectedPairs[candidateIndex].first, score); } } - } - if (residualTopUpBudget > 0) { - std::vector> residualPairs; - residualPairs.reserve(nodeTypeState.residuals.size()); - for (const auto& [nodeId, residual] : nodeTypeState.residuals) { - if (residual <= 0.0 || selectedPPRNodeIds.find(nodeId) != selectedPPRNodeIds.end()) { - continue; - } + addTypedPPRSeedCandidates(baseScores, + baseCandidatesByChannel[channelIndex], + baseNodesAndScores, + extendedMaxScore, + channelIndex, + numEdgeAttrFeatures, + numChannels); + addTypedPPRSeedCandidates(extendedScores, + extendedCandidatesByChannel[channelIndex], + extendedNodesAndScores, + extendedMaxScore, + channelIndex, + numEdgeAttrFeatures, + numChannels); + } - auto scoreIter = scores.find(nodeId); - double pprScore = (scoreIter != scores.end()) ? scoreIter->second : 0.0; - double outputScore = pprScore + residual; - residualPairs.emplace_back(nodeId, outputScore); + std::vector selectedNodes; + std::unordered_set selectedNodeIds; + auto baseSelectedNodes = + selectTypedPPRNodeIds(baseScores, baseCandidatesByChannel, channelQuotas, maxPPRNodes); + selectedNodes.reserve(static_cast(maxPPRNodes)); + selectedNodeIds.reserve(static_cast(maxPPRNodes)); + for (int32_t nodeId : baseSelectedNodes) { + if (static_cast(selectedNodes.size()) >= maxPPRNodes) { + break; } + selectedNodes.push_back(nodeId); + selectedNodeIds.insert(nodeId); + flatIds.push_back(static_cast(nodeId)); + const auto& features = baseScores.at(nodeId); + flatFeatureValues.insert(flatFeatureValues.end(), features.begin(), features.end()); + } - const int32_t residualTopK = std::min(residualTopUpBudget, static_cast(residualPairs.size())); - if (residualTopK > 0) { - // Residual candidates only need selection here; selected - // finalized and residual rows are sorted together below. - if (residualTopK < static_cast(residualPairs.size())) { - std::nth_element(residualPairs.begin(), - residualPairs.begin() + residualTopK, - residualPairs.end(), - higherScore); + if (static_cast(selectedNodes.size()) < maxPPRNodes) { + auto extendedSelectedNodes = + selectTypedPPRNodeIds(extendedScores, extendedCandidatesByChannel, topUpChannelQuotas, maxPPRNodes); + for (int32_t nodeId : extendedSelectedNodes) { + if (static_cast(selectedNodes.size()) >= maxPPRNodes) { + break; } - - for (int32_t rankIdx = 0; rankIdx < residualTopK; ++rankIdx) { - selectedPairs.emplace_back(residualPairs[rankIdx].first, residualPairs[rankIdx].second); + if (selectedNodeIds.find(nodeId) != selectedNodeIds.end()) { + continue; } + selectedNodes.push_back(nodeId); + selectedNodeIds.insert(nodeId); + flatIds.push_back(static_cast(nodeId)); + const auto& features = extendedScores.at(nodeId); + flatFeatureValues.insert(flatFeatureValues.end(), features.begin(), features.end()); } } - if (enableResidualTopUp && selectedPairs.size() > 1) { - std::sort(selectedPairs.begin(), selectedPairs.end(), higherScore); - } - for (const auto& [nodeId, score] : selectedPairs) { - flatIds.push_back(static_cast(nodeId)); - flatWeights.push_back(score); - } - validCounts.push_back(static_cast(selectedPairs.size())); + validCounts.push_back(static_cast(selectedNodes.size())); } - result[nodeTypeId] = {torch::tensor(flatIds, torch::kLong), - torch::tensor(flatWeights, torch::kDouble), - torch::tensor(validCounts, torch::kLong)}; + auto flatWeights = + torch::tensor(flatFeatureValues, torch::kDouble) + .reshape({static_cast(flatIds.size()), static_cast(numEdgeAttrFeatures)}); + result[nodeTypeId] = { + torch::tensor(flatIds, torch::kLong), + flatWeights, + torch::tensor(validCounts, torch::kLong), + }; } + return result; } diff --git a/gigl-core/core/sampling/ppr_forward_push.h b/gigl-core/core/sampling/ppr_forward_push.h index 4698a6998..2b07861c7 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/gigl-core/core/sampling/ppr_forward_push.h @@ -25,6 +25,21 @@ struct SeedNodeTypeState { std::unordered_set queuedNodes; // snapshot captured by drainQueue() }; +// Batched drain result for typed-PPR channels. +struct DrainedTypedPPRQueues { + // Channels that drained queued nodes and need pushResiduals() this iteration. + std::vector activeChannelIndices; + + // Channels that have non-empty uncached frontiers and remaining fetch budget. + 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. + std::unordered_map unionNodesByEdgeTypeId; +}; + // C++ kernel for PPR Forward Push (Andersen et al., 2006). // Hot-loop state lives here; distributed neighbor fetches are driven from Python. // @@ -75,6 +90,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 +137,18 @@ class PPRForwardPush { std::unordered_map> _neighborCache; }; +// Drain several independent PPR states and union their fetch frontier by edge type. +// Used by typed PPR to keep per-channel PPR state separate while issuing one +// shared neighbor fetch for duplicate channel frontier requests. +DrainedTypedPPRQueues drainTypedPPRChannelQueues(const std::vector& states, + const std::vector& fetchIterationCounts, + int32_t maxFetchIterations); + +// Extract and merge typed-PPR channel states in one C++ step. +std::unordered_map> 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..c4c91c294 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -66,6 +66,56 @@ extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, return result; } +static py::tuple drainTypedPPRChannelQueuesWrapper(const py::sequence& states, + const std::vector& fetchIterationCounts, + int32_t maxFetchIterations) { + std::vector statePtrs; + statePtrs.reserve(py::len(states)); + // Sequence iteration and casting touch Python objects, so keep the GIL + // while copying raw C++ state pointers out of the Python container. + for (py::handle stateObj : states) { + statePtrs.push_back(&stateObj.cast()); + } + + // C++ typed drain only reads/mutates PPRForwardPush states and builds C++ + // containers. Reacquire the GIL before constructing the Python tuple. + // REQUIREMENT: no other thread may read or mutate these channel states + // while the GIL is released. The typed sampler drains and pushes each + // channel in a single sequenced loop iteration. + DrainedTypedPPRQueues drained; + { + py::gil_scoped_release release; + drained = drainTypedPPRChannelQueues(statePtrs, fetchIterationCounts, maxFetchIterations); + } + return py::make_tuple(std::move(drained.activeChannelIndices), + std::move(drained.fetchChannelIndices), + std::move(drained.edgeTypeIdsByFetchChannel), + std::move(drained.unionNodesByEdgeTypeId)); +} + +static std::unordered_map> +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. + std::unordered_map> 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 @@ -85,7 +135,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-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index b9f07125b..2bee9780f 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,106 @@ 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 drained = drainTypedPPRChannelQueues(std::vector{&channel0, &channel1}, + /*fetchIterationCounts=*/{0, 0}, + /*maxFetchIterations=*/-1); + + EXPECT_EQ(drained.activeChannelIndices, std::vector({0, 1})); + EXPECT_EQ(drained.fetchChannelIndices, std::vector({0, 1})); + ASSERT_EQ(drained.edgeTypeIdsByFetchChannel.size(), 2); + EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[0], std::vector({0})); + EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[1], std::vector({0})); + + ASSERT_NE(drained.unionNodesByEdgeTypeId.find(0), drained.unionNodesByEdgeTypeId.end()); + EXPECT_EQ(tensorValues(drained.unionNodesByEdgeTypeId.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 drained = drainTypedPPRChannelQueues(std::vector{&channel0, &channel1}, + /*fetchIterationCounts=*/{1, 0}, + /*maxFetchIterations=*/1); + + EXPECT_EQ(drained.activeChannelIndices, std::vector({0, 1})); + EXPECT_EQ(drained.fetchChannelIndices, std::vector({1})); + ASSERT_EQ(drained.edgeTypeIdsByFetchChannel.size(), 1); + EXPECT_EQ(drained.edgeTypeIdsByFetchChannel[0], std::vector({0})); + + ASSERT_NE(drained.unionNodesByEdgeTypeId.find(0), drained.unionNodesByEdgeTypeId.end()); + EXPECT_EQ(tensorValues(drained.unionNodesByEdgeTypeId.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, 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. diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 78de00436..5bedae5a4 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,13 @@ from graphlearn_torch.utils import merge_dict from gigl.distributed.base_sampler import BaseDistNeighborSampler +from gigl.distributed.utils.distributed_typed_sampler import ( + HeteroPPRResult, + PPRResult, + TypedPPRChannelKey, + build_edge_type_channel_group_edge_type_ids, + parse_typed_channel_quota_groups, +) 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 +70,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) @@ -72,7 +94,12 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): **Heterogeneous (HeteroData)** — one PPR edge type per ``(seed_type, neighbor_type)`` pair, with ``"ppr"`` as the relation: - ``data[(seed_type, "ppr", neighbor_type)].edge_index``: same format as above. - - ``data[(seed_type, "ppr", neighbor_type)].edge_attr``: same format as above. + - ``data[(seed_type, "ppr", neighbor_type)].edge_attr``: scalar PPR + score for regular PPR. For typed PPR, edge attrs are multi-column: + ``[best_calibrated_score, calibrated_channel_scores..., channel_presence_bits...]``. + Typed-PPR scores are calibrated within each channel/seed pool and + globally ranked by the best calibrated score. Channel columns follow + the insertion order of ``typed_channel_quotas``. Args: alpha: Restart probability (teleport probability back to seed). Higher values @@ -89,6 +116,21 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): during Forward Push when fewer than ``max_ppr_nodes`` finalized PPR scores are available. num_neighbors_per_hop: Maximum number of neighbors to fetch per hop. + typed_channel_quotas: Optional top-k quotas for typed PPR traversal + channels. Each channel may contribute up to its quota to the + candidate pool; the final returned sequence is still capped by + ``max_ppr_nodes``. Quotas may sum above ``max_ppr_nodes`` to give + sparse or overlapping channels room to fill the sequence. + Example:: + + typed_channel_quotas = { + ("user", "views", "item"): 64, + ( + ("user", "likes", "item"), + ("user", "shares", "item"), + ): 32, + } + degree_tensors: Pre-computed total-degree tensors (int32). Homogeneous graphs use a single tensor; heterogeneous graphs use tensors keyed by NodeType. The colocated and graph-store loader paths retrieve @@ -106,6 +148,7 @@ def __init__( num_neighbors_per_hop: int = 100_000, degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], max_fetch_iterations: Optional[int] = None, + typed_channel_quotas: Optional[dict[TypedPPRChannelKey, int]] = None, **kwargs, ): super().__init__(*args, **kwargs) @@ -147,6 +190,15 @@ def __init__( ] self._is_homogeneous = True + typed_channel_groups, self._typed_ppr_channel_quotas = ( + parse_typed_channel_quota_groups(typed_channel_quotas) + ) + if self._typed_ppr_channel_quotas is not None: + if self._is_homogeneous: + raise ValueError( + "Typed PPR channel quotas are only supported for heterogeneous PPR sampling." + ) + # 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 +261,18 @@ def __init__( for node_type in all_node_types ] + if typed_channel_groups is not None: + self._typed_ppr_channel_to_node_type_id_to_edge_type_ids = ( + build_edge_type_channel_group_edge_type_ids( + edge_type_groups=typed_channel_groups, + edge_type_to_edge_type_id=self._etype_to_etype_id, + node_type_to_edge_types=self._node_type_to_edge_types, + node_types=self._ntype_id_to_ntype, + ) + ) + else: + self._typed_ppr_channel_to_node_type_id_to_edge_type_ids = [] + def _convert_degree_tensors_to_dict( self, degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], @@ -304,12 +368,8 @@ 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]], - ]: + ppr_node_limit: int, + ) -> PPRResult: """Extract PPR neighbors from a completed C++ Forward Push state. The C++ kernel indexes node types by compact integer IDs for speed. @@ -317,7 +377,7 @@ def _extract_ppr_state_top_k( preserves the homogeneous return shape expected by the rest of the sampler. - ``max_ppr_nodes`` is the maximum number of nodes returned for each + ``ppr_node_limit`` is the maximum number of nodes returned for each source node. If residual top-up is enabled, residual candidates count against this cap. @@ -334,7 +394,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, + ppr_node_limit, self._enable_residual_topup, ) @@ -367,15 +427,42 @@ def _extract_ppr_state_top_k( node_type_to_valid_counts, ) + def _extract_typed_ppr_state_top_k( + self, + ppr_states, + 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, + 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,14 +560,12 @@ 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, @@ -489,6 +574,115 @@ async def _compute_ppr_scores( self._max_ppr_nodes, ) + async def _compute_typed_ppr_scores( + self, + seed_nodes: torch.Tensor, + seed_node_type: NodeType, + ) -> 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``. + + Returns: + Heterogeneous PPR extraction output with typed edge-attribute + feature vectors. + """ + assert self._typed_ppr_channel_quotas is not None + channel_quotas = self._typed_ppr_channel_quotas + 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(): + 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: + ( + active_channel_indices, + fetch_channel_indices, + edge_type_ids_by_fetch_channel, + union_nodes_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 active_channel_indices: + break + + fetched_by_channel: list[ + dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + ] = [dict() for _ in ppr_states] + + if union_nodes_by_edge_type_id: + union_fetched_by_edge_type_id = await self._batch_fetch_neighbors( + union_nodes_by_edge_type_id + ) + for channel_index, edge_type_ids in zip( + fetch_channel_indices, + edge_type_ids_by_fetch_channel, + ): + 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 + if edge_type_id in union_fetched_by_edge_type_id + } + + # 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 active_channel_indices + ] + if push_tasks: + await asyncio.gather(*push_tasks) + + return await loop.run_in_executor( + None, + self._extract_typed_ppr_state_top_k, + ppr_states, + channel_quotas, + device, + ) + async def _sample_from_nodes( self, inputs: NodeSamplerInput, @@ -590,15 +784,26 @@ 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 - ] - ) + if self._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, + ) + 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 +819,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 +874,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/sampler_options.py b/gigl/distributed/sampler_options.py index b093b8e10..55044c836 100644 --- a/gigl/distributed/sampler_options.py +++ b/gigl/distributed/sampler_options.py @@ -16,6 +16,16 @@ logger = Logger() +TypedPPRChannelKey = Union[EdgeType, tuple[EdgeType, ...]] +"""A 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. +""" + @dataclass(frozen=True) class KHopNeighborSamplerOptions: @@ -41,6 +51,8 @@ class PPRSamplerOptions: - ``edge_index``: ``[2, N]`` int64 — row 0 is local seed indices, row 1 is local neighbor indices. - ``edge_attr``: ``[N]`` float — PPR score for each (seed, neighbor) pair. + Typed PPR emits multi-column edge attrs: + ``[best_score, channel_scores..., channel_presence_bits...]``. For homogeneous graphs these live directly on ``data.edge_index`` / ``data.edge_attr``. @@ -75,6 +87,33 @@ class PPRSamplerOptions: The algorithm still runs to convergence — re-enqueued nodes propagate through cached neighbors at negligible cost. ``None`` (default) means no fetch limit. + typed_channel_quotas: Optional top-k quotas for typed PPR + traversal channels defined by canonical edge-type allowlists. Keys + may be either a single canonical edge type + ``(src_type, relation, dst_type)`` or a tuple of canonical edge + types. Each key defines one traversal channel whose PPR state may + traverse only those exact edge types. + Channel order follows the insertion order of this mapping. Each + channel may contribute up to its quota to the candidate pool, while + the final returned sequence remains capped by ``max_ppr_nodes``. + Quotas may sum above ``max_ppr_nodes`` to give sparse or + overlapping channels room to fill the sequence. + Example:: + + typed_channel_quotas = { + ("user", "views", "item"): 64, + ( + ("user", "likes", "item"), + ("user", "shares", "item"), + ): 32, + } + + If residual top-up is enabled and the base merge emits fewer than + ``max_ppr_nodes``, the sampler appends discovered-but-unpushed + residual candidates from the same completed PPR states. Those + top-up candidates are deduplicated, globally ranked, and emitted on + the same mass scale as finalized PPR scores: + ``ppr_score + residual``. """ alpha: float = 0.5 @@ -83,6 +122,7 @@ class PPRSamplerOptions: enable_residual_topup: bool = True num_neighbors_per_hop: int = 1_000 max_fetch_iterations: Optional[int] = None + typed_channel_quotas: Optional[dict[TypedPPRChannelKey, int]] = None SamplerOptions = Union[KHopNeighborSamplerOptions, PPRSamplerOptions] diff --git a/gigl/distributed/utils/dist_sampler.py b/gigl/distributed/utils/dist_sampler.py index e0289b26c..79ed2b34e 100644 --- a/gigl/distributed/utils/dist_sampler.py +++ b/gigl/distributed/utils/dist_sampler.py @@ -85,6 +85,7 @@ def create_dist_sampler( enable_residual_topup=sampler_options.enable_residual_topup, max_fetch_iterations=sampler_options.max_fetch_iterations, num_neighbors_per_hop=sampler_options.num_neighbors_per_hop, + typed_channel_quotas=sampler_options.typed_channel_quotas, degree_tensors=degree_tensors, ) else: diff --git a/gigl/distributed/utils/distributed_typed_sampler.py b/gigl/distributed/utils/distributed_typed_sampler.py new file mode 100644 index 000000000..0dd4134a3 --- /dev/null +++ b/gigl/distributed/utils/distributed_typed_sampler.py @@ -0,0 +1,156 @@ +"""Typed-PPR helpers for distributed sampler implementations.""" + +from typing import Optional, Sequence, 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, ...]] +# Parsed typed-channel edge-type allowlists, ordered to match quota order. +TypedPPRChannelEdgeTypeGroups = list[tuple[EdgeType, ...]] + + +def parse_typed_channel_quota_groups( + typed_channel_quotas: Optional[dict[TypedPPRChannelKey, int]], +) -> tuple[Optional[TypedPPRChannelEdgeTypeGroups], Optional[list[int]]]: + """Validate 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. + + 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 quota is not a positive integer, or 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] = [] + invalid_quotas: dict[TypedPPRChannelKey, object] = {} + + def is_canonical_edge_type(value: object) -> bool: + """Return whether ``value`` has GraphLearn's canonical EdgeType 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 not isinstance(quota, int) or isinstance(quota, bool) or quota <= 0: + invalid_quotas[edge_type_key] = quota + continue + 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) + + if invalid_quotas: + raise ValueError( + "typed_channel_quotas must contain only positive integer quotas, " + f"got {invalid_quotas}." + ) + 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], +) -> list[list[list[int]]]: + """Convert typed-channel edge-type allowlists to PPRForwardPush IDs. + + Returns one traversal map per typed channel. Each traversal map is indexed + as ``[node_type_id][allowed_edge_type_ids]`` and tells the C++ forward push + state which edge-type IDs may be traversed from each node type for that + channel. + + 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: + A list of traversal maps aligned with ``edge_type_groups``. Each + traversal map is indexed by integer node-type ID and stores the allowed + integer edge-type IDs for 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: list[list[list[int]]] = [] + 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}. Edge types must exist in the " + "graph and must not be label edge types." + ) + + channel_edge_type_set = set(channel_edge_types) + node_type_id_to_channel_edge_type_ids = [ + [ + edge_type_to_edge_type_id[edge_type] + for edge_type in node_type_to_edge_types.get(node_type, []) + if edge_type in channel_edge_type_set + ] + for node_type in node_types + ] + 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_ppr_sampler_test.py b/tests/unit/distributed/dist_ppr_sampler_test.py index 604e650c6..26b3aec57 100644 --- a/tests/unit/distributed/dist_ppr_sampler_test.py +++ b/tests/unit/distributed/dist_ppr_sampler_test.py @@ -41,6 +41,10 @@ from gigl.distributed.dist_ablp_neighborloader import DistABLPLoader from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.sampler_options import PPRSamplerOptions +from gigl.distributed.utils.distributed_typed_sampler import ( + build_edge_type_channel_group_edge_type_ids, + parse_typed_channel_quota_groups, +) from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, @@ -822,6 +826,64 @@ def test_ppr_sampler_homogeneous_ablp(self) -> None: """Verify PPR handles homogeneous ABLP seed dictionaries.""" mp.spawn(fn=_run_ppr_labeled_homogeneous_ablp_loader_check, args=()) + 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, "positive quotas"): + parse_typed_channel_quota_groups({USER_TO_STORY: 0}) + 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()