Skip to content

Refactor PPR top-k selection helper #723

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mkolodner-sc wants to merge 2 commits into
base: main
Choose a base branch
from
+126 −78
Open

Refactor PPR top-k selection helper #723

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
930569f
Extract PPR top-k selection helper
mkolodner-sc Jul 28, 2026
d06fbf1
Address PPR helper review comments
mkolodner-sc Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
204 changes: 126 additions & 78 deletions gigl-core/core/sampling/ppr_forward_push.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,11 +266,119 @@ void PPRForwardPush::pushResiduals(
}
}

// Helper function for selecting one seed/node-type's finalized PPR rows.
//
// Inputs:
// nodeTypeState: finalized PPR scores and residuals for one seed/node type.
// finalizedPPRNodeLimit: maximum finalized-PPR rows to select before top-up.
//
// Expected output: (node_id, raw_ppr_score) pairs selected by raw PPR score.
// The order is unspecified; callers that emit these directly should sort the
// returned vector before writing output tensors.
static std::vector<std::pair<int32_t, double>> selectFinalizedPPRPairs(const SeedNodeTypeState& nodeTypeState,
int32_t finalizedPPRNodeLimit) {
const auto& scores = nodeTypeState.pprScores;

const int32_t numReturnedPairs = std::min(finalizedPPRNodeLimit, static_cast<int32_t>(scores.size()));
std::vector<std::pair<int32_t, double>> selectedPairs;
selectedPairs.reserve(static_cast<size_t>(numReturnedPairs));
if (numReturnedPairs > 0) {
std::vector<std::pair<int32_t, double>> scorePairs(scores.begin(), scores.end());
if (numReturnedPairs < static_cast<int32_t>(scorePairs.size())) {
std::nth_element(scorePairs.begin(),
scorePairs.begin() + numReturnedPairs,
scorePairs.end(),
[](const auto& a, const auto& b) { return a.second > b.second; });
}

for (int32_t rankIdx = 0; rankIdx < numReturnedPairs; ++rankIdx) {
selectedPairs.emplace_back(scorePairs[rankIdx].first, scorePairs[rankIdx].second);
}
}

return selectedPairs;
}

// Helper function for extending one seed/node-type's selected PPR rows with
// residual top-up candidates.
//
// Inputs:
// nodeTypeState: finalized PPR scores and residuals for one seed/node type.
// selectedPairs: mutable finalized-PPR rows already selected for this seed.
// sequenceLength: maximum total rows after residual top-up.
//
// Expected output: selectedPairs has up to sequenceLength rows after appending
// highest-scoring residual candidates that are not already selected. This helper
// does not sort selectedPairs; callers sort only when their output needs it.
static void appendResidualTopUpPairs(const SeedNodeTypeState& nodeTypeState,
std::vector<std::pair<int32_t, double>>& selectedPairs,
int32_t sequenceLength) {
const int32_t residualTopUpBudget =
std::max<int32_t>(0, sequenceLength - static_cast<int32_t>(selectedPairs.size()));
if (residualTopUpBudget > 0) {
const std::unordered_map<int32_t, double>& pprScoresByNodeId = nodeTypeState.pprScores;
std::unordered_set<int32_t> selectedPPRNodeIds;
selectedPPRNodeIds.reserve(selectedPairs.size());
for (const auto& selectedPair : selectedPairs) {
selectedPPRNodeIds.insert(selectedPair.first);
}

std::vector<std::pair<int32_t, double>> residualPairs;
residualPairs.reserve(nodeTypeState.residuals.size());
for (const auto& [nodeId, residual] : nodeTypeState.residuals) {
// Forward push residuals are non-negative in normal operation. Pushed
// nodes remain in the map with zero residual, so skip drained entries
// and any unexpected non-positive values.
if (residual <= 0.0 || selectedPPRNodeIds.find(nodeId) != selectedPPRNodeIds.end()) {
Comment thread
mkolodner-sc marked this conversation as resolved.
continue;
}

std::unordered_map<int32_t, double>::const_iterator pprScoreIter = pprScoresByNodeId.find(nodeId);
double pprScore = (pprScoreIter != pprScoresByNodeId.end()) ? pprScoreIter->second : 0.0;
double outputScore = pprScore + residual;
residualPairs.emplace_back(nodeId, outputScore);
}

const int32_t residualTopK = std::min(residualTopUpBudget, static_cast<int32_t>(residualPairs.size()));
if (residualTopK > 0) {
if (residualTopK < static_cast<int32_t>(residualPairs.size())) {
std::nth_element(residualPairs.begin(),
residualPairs.begin() + residualTopK,
residualPairs.end(),
[](const auto& a, const auto& b) { return a.second > b.second; });
}

for (int32_t rankIdx = 0; rankIdx < residualTopK; ++rankIdx) {
selectedPairs.emplace_back(residualPairs[rankIdx].first, residualPairs[rankIdx].second);
}
}
}
}

// Helper function for moving finalized-PPR rows onto the same score scale as
// residual top-up rows.
//
// Inputs:
// nodeTypeState: residual table for one seed/node type.
// selectedPairs: mutable finalized-PPR rows with raw PPR scores.
//
// Expected output: selectedPairs scores are updated in-place to
// ppr_score + residual(node) when residual mass exists for that node.
static void addResidualMassToPPRPairs(const SeedNodeTypeState& nodeTypeState,
std::vector<std::pair<int32_t, double>>& selectedPairs) {
for (auto& [nodeId, score] : selectedPairs) {
auto residualIter = nodeTypeState.residuals.find(nodeId);
if (residualIter != nodeTypeState.residuals.end()) {
score += residualIter->second;
}
}
}

std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> PPRForwardPush::
extractTopKWithResidualTopUp(int32_t maxPPRNodes, bool enableResidualTopUp) {
TORCH_CHECK(maxPPRNodes >= 0, "maxPPRNodes must be non-negative, got ", maxPPRNodes, ".");

std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> result;
std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> 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.
Expand All @@ -281,95 +389,35 @@ std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tens

for (int32_t seedIdx = 0; seedIdx < _batchSize; ++seedIdx) {
const auto& nodeTypeState = _state[seedIdx][nodeTypeId];
const auto& scores = nodeTypeState.pprScores;
const auto higherScore = [](const auto& a, const auto& b) { return a.second > b.second; };

const int32_t topK = std::min(maxPPRNodes, static_cast<int32_t>(scores.size()));
const int32_t residualTopUpBudget = enableResidualTopUp ? maxPPRNodes - topK : 0;
std::vector<std::pair<int32_t, double>> selectedPairs;
selectedPairs.reserve(static_cast<size_t>(topK) + static_cast<size_t>(residualTopUpBudget));
std::unordered_set<int32_t> selectedPPRNodeIds;
if (topK > 0) {
if (residualTopUpBudget > 0) {
selectedPPRNodeIds.reserve(static_cast<size_t>(topK));
}
std::vector<std::pair<int32_t, double>> 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<int32_t>(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 (enableResidualTopUp) {
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);
}
}
auto selectedPairs = selectFinalizedPPRPairs(nodeTypeState, maxPPRNodes);
if (enableResidualTopUp) {
addResidualMassToPPRPairs(nodeTypeState, selectedPairs);
appendResidualTopUpPairs(nodeTypeState, selectedPairs, maxPPRNodes);
}

if (residualTopUpBudget > 0) {
std::vector<std::pair<int32_t, double>> 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<int32_t>(residualPairs.size()));
if (residualTopK > 0) {
// Residual candidates only need selection here; selected
// finalized and residual rows are sorted together below.
if (residualTopK < static_cast<int32_t>(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);
}
}
// The selection helpers use nth_element, which selects the right rows
// but does not order them. Sort the selected rows once to preserve the
// emitted ordering contract. With residual top-up enabled, this matches
// the previous behavior: selected finalized and top-up rows are ordered
// together by emitted score, so top-up rows may interleave after selection.
if (selectedPairs.size() > 1) {
std::sort(selectedPairs.begin(), selectedPairs.end(), [](const auto& a, const auto& b) {
return a.second > b.second;
});
}

if (enableResidualTopUp && selectedPairs.size() > 1) {
std::sort(selectedPairs.begin(), selectedPairs.end(), higherScore);
}
for (const auto& [nodeId, score] : selectedPairs) {
flatIds.push_back(static_cast<int64_t>(nodeId));
flatWeights.push_back(score);
}
validCounts.push_back(static_cast<int64_t>(selectedPairs.size()));
}

result[nodeTypeId] = {torch::tensor(flatIds, torch::kLong),
torch::tensor(flatWeights, torch::kDouble),
torch::tensor(validCounts, torch::kLong)};
extractedPPRByNodeTypeId[nodeTypeId] = {torch::tensor(flatIds, torch::kLong),
torch::tensor(flatWeights, torch::kDouble),
torch::tensor(validCounts, torch::kLong)};
}
return result;
return extractedPPRByNodeTypeId;
}

int32_t PPRForwardPush::getTotalDegree(int32_t nodeId, int32_t nodeTypeId) const {
Expand Down