Skip to content

Localize PersonalizedPageRank traversal #6

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
vndee wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Localize PersonalizedPageRank traversal #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ea663a6
Localize PersonalizedPageRank traversal
vndee Apr 6, 2026
e392a16
Address PR review feedback
vndee Apr 7, 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
200 changes: 153 additions & 47 deletions internal/graph/graph.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package graph

import (
"context"
"math"
"sort"
"sync"
"time"
)
Expand Down Expand Up @@ -178,8 +180,8 @@ type SubgraphEdge struct {

// SubgraphResult holds the BFS neighborhood with both nodes and edges.
type SubgraphResult struct {
Nodes map[string]int // entityID -> hop distance from seed
Edges []SubgraphEdge // edges connecting nodes within the subgraph
Nodes map[string]int // entityID -> hop distance from seed
Edges []SubgraphEdge // edges connecting nodes within the subgraph
}

// Subgraph returns the N-hop ego-graph around seeds with full edge data.
Expand Down Expand Up @@ -389,70 +391,94 @@ func (g *Graph) WeightedNeighbors(seeds []string, maxHops int, minWeight float64
// --- Personalized PageRank ---

// PersonalizedPageRank computes PPR scores seeded from the given entity IDs.
// alpha is the teleport probability (typically 0.15), maxIter limits iterations,
// and epsilon is the convergence threshold.
func (g *Graph) PersonalizedPageRank(seeds []string, alpha float64, maxIter int, epsilon float64) map[string]float64 {
g.mu.RLock()
defer g.mu.RUnlock()

// The random walk is restricted to the local neighborhood within maxHops of the
// seeds so per-query work scales with the relevant subgraph rather than the
// entire KB graph. alpha is the teleport probability (typically 0.15), maxIter
// limits iterations, and epsilon is the convergence threshold.
func (g *Graph) PersonalizedPageRank(
ctx context.Context,
seeds []string,
maxHops int,
alpha float64,
maxIter int,
epsilon float64,
) (map[string]float64, error) {
if len(seeds) == 0 {
return nil
return nil, nil
}

// Build the personalization vector (uniform over seeds).
seedWeight := 1.0 / float64(len(seeds))
personal := make(map[string]float64, len(seeds))
for _, s := range seeds {
personal[s] = seedWeight
if err := ctx.Err(); err != nil {
return nil, err
}
Comment on lines 393 to 411

@coderabbitai coderabbitai Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Wrap context errors with descriptive context.

Per coding guidelines, errors should be wrapped using fmt.Errorf("context: %w", err). The raw ctx.Err() returns at lines 409 and similar locations throughout this function lack context about where the cancellation occurred.

Proposed fix 
 	if err := ctx.Err(); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("ppr: context check: %w", err)
 	}

Apply similar wrapping to other ctx.Err() returns at lines 446-448, 456-458, 472-474, and in localPPRNeighborhoodLocked at lines 512-514.

As per coding guidelines: **/*.go: Always wrap errors with context using fmt.Errorf("context: %w", err). Never return naked errors.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// PersonalizedPageRank computes PPR scores seeded from the given entity IDs.
// alpha is the teleport probability (typically 0.15), maxIter limits iterations,
// and epsilon is the convergence threshold.
func (g *Graph) PersonalizedPageRank(seeds []string, alpha float64, maxIter int, epsilon float64) map[string]float64 {
g.mu.RLock()
defer g.mu.RUnlock()
// The random walk is restricted to the local neighborhood within maxHops of the
// seeds so per-query work scales with the relevant subgraph rather than the
// entire KB graph. alpha is the teleport probability (typically 0.15), maxIter
// limits iterations, and epsilon is the convergence threshold.
func (g *Graph) PersonalizedPageRank(
ctx context.Context,
seeds []string,
maxHops int,
alpha float64,
maxIter int,
epsilon float64,
) (map[string]float64, error) {
if len(seeds) == 0 {
return nil
return nil, nil
}
// Build the personalization vector (uniform over seeds).
seedWeight := 1.0 / float64(len(seeds))
personal := make(map[string]float64, len(seeds))
for _, s := range seeds {
personal[s] = seedWeight
if err := ctx.Err(); err != nil {
return nil, err
}
// PersonalizedPageRank computes PPR scores seeded from the given entity IDs.
// The random walk is restricted to the local neighborhood within maxHops of the
// seeds so per-query work scales with the relevant subgraph rather than the
// entire KB graph. alpha is the teleport probability (typically 0.15), maxIter
// limits iterations, and epsilon is the convergence threshold.
func (g *Graph) PersonalizedPageRank(
ctx context.Context,
seeds []string,
maxHops int,
alpha float64,
maxIter int,
epsilon float64,
) (map[string]float64, error) {
if len(seeds) == 0 {
return nil, nil
}
if err := ctx.Err(); err != nil {
return nil, fmt.Errorf("ppr: context check: %w", err)
}
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@internal/graph/graph.go` around lines 392 - 410, The raw ctx.Err() returns in
PersonalizedPageRank and localPPRNeighborhoodLocked should be wrapped with
contextual error messages; replace direct returns of ctx.Err() (e.g., in
PersonalizedPageRank around the initial check and the later checks previously at
lines noted in the review, and in localPPRNeighborhoodLocked) with
fmt.Errorf("context: %w", err) so callers get descriptive context; search for
ctx.Err() usages inside the PersonalizedPageRank method and
localPPRNeighborhoodLocked helper and update each return to wrap the error using
fmt.Errorf("context: %w", err).


// Collect all nodes reachable in the graph (union of forward + reverse keys).
allNodes := make(map[string]struct{})
for id := range g.forward {
allNodes[id] = struct{}{}
g.mu.RLock()
nodes, neighbors, err := g.localPPRNeighborhoodLocked(ctx, seeds, maxHops)
g.mu.RUnlock()
if err != nil {
return nil, err
}
for id := range g.reverse {
allNodes[id] = struct{}{}
if len(nodes) == 0 {
return nil, nil
}
Comment on lines +413 to 421

@coderabbitai coderabbitai Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Wrap error from localPPRNeighborhoodLocked.

The error returned from the helper should be wrapped with context to aid debugging.

Proposed fix 
 	g.mu.RLock()
 	nodes, neighbors, err := g.localPPRNeighborhoodLocked(ctx, seeds, maxHops)
 	g.mu.RUnlock()
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("ppr neighborhood: %w", err)
 	}

As per coding guidelines: **/*.go: Always wrap errors with context using fmt.Errorf("context: %w", err). Never return naked errors.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
g.mu.RLock()
nodes, neighbors, err := g.localPPRNeighborhoodLocked(ctx, seeds, maxHops)
g.mu.RUnlock()
if err != nil {
return nil, err
}
for id := range g.reverse {
allNodes[id] = struct{}{}
if len(nodes) == 0 {
return nil, nil
}
g.mu.RLock()
nodes, neighbors, err := g.localPPRNeighborhoodLocked(ctx, seeds, maxHops)
g.mu.RUnlock()
if err != nil {
return nil, fmt.Errorf("ppr neighborhood: %w", err)
}
if len(nodes) == 0 {
return nil, nil
}
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@internal/graph/graph.go` around lines 412 - 420, The call to
g.localPPRNeighborhoodLocked returns an unwrapped error; update the error return
to wrap it with context using fmt.Errorf so callers get useful debug info (e.g.,
replace "return nil, err" with "return nil,
fmt.Errorf(\"localPPRNeighborhoodLocked failed: %w\", err)"). Ensure you import
fmt if needed and keep the surrounding lock/unlock and nil/empty-nodes behavior
unchanged; references: g.localPPRNeighborhoodLocked and the error-return site in
the function that holds g.mu.RLock/g.mu.RUnlock.


// Degree = total edges (forward + reverse) for undirected view.
degree := make(map[string]int, len(allNodes))
for id := range allNodes {
degree[id] = len(g.forward[id]) + len(g.reverse[id])
personal := make(map[string]float64, len(seeds))
for _, s := range seeds {
if _, ok := neighbors[s]; !ok {
continue
}
personal[s]++
}

// Initialize ranks.
n := float64(len(allNodes))
rank := make(map[string]float64, len(allNodes))
for id := range allNodes {
rank[id] = 1.0 / n
if len(personal) == 0 {
return nil, nil
}
seedWeight := 1.0 / float64(len(personal))
for id := range personal {
personal[id] = seedWeight
}

newRank := make(map[string]float64, len(allNodes))
rank := make(map[string]float64, len(nodes))
newRank := make(map[string]float64, len(nodes))
initialRank := 1.0 / float64(len(nodes))
for _, id := range nodes {
rank[id] = initialRank
newRank[id] = 0
}

for iter := 0; iter < maxIter; iter++ {
// Reset newRank for this iteration (reuse allocation).
for id := range newRank {
delete(newRank, id)
if err := ctx.Err(); err != nil {
return nil, err
}

// Distribute rank from each node to neighbors.
for id := range allNodes {
if degree[id] == 0 {
continue
for _, id := range nodes {
newRank[id] = 0
}

for i, id := range nodes {
if i%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, err
}
}
share := rank[id] / float64(degree[id])
for _, e := range g.forward[id] {
newRank[e.TargetID] += share
if len(neighbors[id]) == 0 {
continue
}
for _, e := range g.reverse[id] {
newRank[e.TargetID] += share
share := rank[id] / float64(len(neighbors[id]))
for j, nextID := range neighbors[id] {
Comment on lines +461 to +465

@coderabbitai coderabbitai Bot Apr 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Localized sinks currently leak probability mass.

Restricting the walk to the local subgraph can create sinks, e.g. maxHops == 0 or an isolated seed. The current continue path drops that mass, so a seed-only neighborhood converges to alpha instead of 1.0. Redistribute dangling mass through personal (or add self-loops) before the update step; a small maxHops == 0 regression test would catch this.

🧮 Proposed fix 
+		danglingMass := 0.0
 		for i, id := range nodes {
 			if i%256 == 0 {
 				if err := ctx.Err(); err != nil {
 					return nil, err
 				}
 			}
 			if len(neighbors[id]) == 0 {
+				danglingMass += rank[id]
 				continue
 			}
 			share := rank[id] / float64(len(neighbors[id]))
@@
-			nr := alpha*personal[id] + (1-alpha)*newRank[id]
+			nr := alpha*personal[id] + (1-alpha)*(newRank[id]+danglingMass*personal[id])

Also applies to: 482-483

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@internal/graph/graph.go` around lines 461 - 465, The loop drops probability
mass when a node has no neighbors (e.g., isolated seed or maxHops==0); instead
of "continue", detect dangling nodes (len(neighbors[id])==0) and redistribute
that node's rank into the personalization vector (personal[id] += rank[id]) or
treat it as a self-loop (compute share = rank[id] and add to nextRank[id])
before the neighbor-distribution step so total mass is conserved; apply the same
fix at the other occurrence around the code that currently checks
len(neighbors[id]) (the noted spots at the later block) to ensure no localized
sinks leak probability.

if j%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, err
}
}
newRank[nextID] += share
}
}

// Apply teleport.
maxDiff := 0.0
for id := range allNodes {
for i, id := range nodes {
if i%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, err
}
}
nr := alpha*personal[id] + (1-alpha)*newRank[id]
diff := math.Abs(nr - rank[id])
if diff > maxDiff {
Expand All @@ -466,7 +492,87 @@ func (g *Graph) PersonalizedPageRank(seeds []string, alpha float64, maxIter int,
}
}

return rank
return rank, nil
}

func (g *Graph) localPPRNeighborhoodLocked(
ctx context.Context,
seeds []string,
maxHops int,
) ([]string, map[string][]string, error) {
nodes := make(map[string]struct{}, len(seeds))
dist := make(map[string]int, len(seeds))
queue := make([]string, 0, len(seeds))

for _, id := range seeds {
if _, seen := dist[id]; seen {
continue
}
nodes[id] = struct{}{}
dist[id] = 0
queue = append(queue, id)
}

for head := 0; head < len(queue); head++ {
if head%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, nil, err
}
}
cur := queue[head]
hop := dist[cur]
if hop >= maxHops {
continue
}
for _, edges := range [][]Edge{g.forward[cur], g.reverse[cur]} {
for edgeIdx, e := range edges {
if edgeIdx%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, nil, err
}
}
nodes[e.TargetID] = struct{}{}
if _, seen := dist[e.TargetID]; seen {
continue
Comment on lines +527 to +536

Copilot AI Apr 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In localPPRNeighborhoodLocked, cancellation is only checked every 256 queue pops, but the nested edge scan for a single high-degree node can still take a long time after the request is canceled. Consider adding a periodic ctx.Err() check inside the inner edge loops so cancellation latency is bounded by edge-scan work, not just queue progress.

Copilot uses AI. Check for mistakes.

@vndee vndee Apr 7, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in e392a16. localPPRNeighborhoodLocked now checks ctx.Err() inside the BFS edge-scan loop so cancellation latency is bounded by edge-scan work, not just queue progress.

}
dist[e.TargetID] = hop + 1
queue = append(queue, e.TargetID)
}
}
}

nodeIDs := make([]string, 0, len(nodes))
neighbors := make(map[string][]string, len(nodes))
for id := range nodes {
nodeIDs = append(nodeIDs, id)
}
Comment on lines +544 to +548

Copilot AI Apr 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

localPPRNeighborhoodLocked builds nodeIDs by iterating a map (for id := range nodes), which makes the subsequent PageRank iterations run over nodes in nondeterministic order. Because rank propagation uses floating-point accumulation, this can lead to run-to-run score jitter and potentially flaky tests. Consider sorting nodeIDs (e.g., sort.Strings) before returning/using it so PPR traversal order is deterministic.

Copilot uses AI. Check for mistakes.
sort.Strings(nodeIDs)
for i, id := range nodeIDs {
if i%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, nil, err
}
}
localNeighbors := make([]string, 0, len(g.forward[id])+len(g.reverse[id]))
edgeCount := 0
for _, edges := range [][]Edge{g.forward[id], g.reverse[id]} {
for _, e := range edges {
if edgeCount%256 == 0 {
if err := ctx.Err(); err != nil {
return nil, nil, err
}
}
edgeCount++
if _, ok := nodes[e.TargetID]; !ok {
continue
}
localNeighbors = append(localNeighbors, e.TargetID)
}
}
neighbors[id] = localNeighbors
}

return nodeIDs, neighbors, nil
}

// --- Helpers ---
Expand Down
53 changes: 50 additions & 3 deletions internal/graph/graph_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package graph

import (
"context"
"errors"
"testing"
"time"
)
Expand Down Expand Up @@ -333,7 +335,10 @@ func TestPersonalizedPageRank_ConvergesOnSeeds(t *testing.T) {
addEdge(g, "a", "b", "r1", "knows", 1.0)
addEdge(g, "b", "c", "r2", "knows", 1.0)

ranks := g.PersonalizedPageRank([]string{"a"}, 0.15, 50, 1e-8)
ranks, err := g.PersonalizedPageRank(context.Background(), []string{"a"}, 2, 0.15, 50, 1e-8)
if err != nil {
t.Fatal(err)
}

if ranks["a"] <= ranks["c"] {
t.Errorf("seed 'a' should rank higher than distant 'c': a=%.6f c=%.6f", ranks["a"], ranks["c"])
Expand All @@ -349,7 +354,10 @@ func TestPersonalizedPageRank_HubsRankHigher(t *testing.T) {
addEdge(g, "h", "d", "r4", "knows", 1.0)
addEdge(g, "e", "d", "r5", "knows", 1.0)

ranks := g.PersonalizedPageRank([]string{"a"}, 0.15, 50, 1e-8)
ranks, err := g.PersonalizedPageRank(context.Background(), []string{"a"}, 3, 0.15, 50, 1e-8)
if err != nil {
t.Fatal(err)
}

if ranks["h"] <= ranks["e"] {
t.Errorf("hub 'h' should rank higher than leaf 'e': h=%.6f e=%.6f", ranks["h"], ranks["e"])
Expand All @@ -360,8 +368,47 @@ func TestPersonalizedPageRank_EmptySeeds(t *testing.T) {
g := New()
addEdge(g, "a", "b", "r1", "knows", 1.0)

ranks := g.PersonalizedPageRank(nil, 0.15, 20, 1e-6)
ranks, err := g.PersonalizedPageRank(context.Background(), nil, 2, 0.15, 20, 1e-6)
if err != nil {
t.Fatal(err)
}
if ranks != nil {
t.Errorf("empty seeds should return nil, got %v", ranks)
}
}

func TestPersonalizedPageRank_LocalNeighborhoodExcludesFarNodes(t *testing.T) {
g := New()
addEdge(g, "a", "b", "r1", "knows", 1.0)
addEdge(g, "b", "c", "r2", "knows", 1.0)
addEdge(g, "x", "y", "r3", "knows", 1.0)
addEdge(g, "y", "z", "r4", "knows", 1.0)

ranks, err := g.PersonalizedPageRank(context.Background(), []string{"a"}, 1, 0.15, 50, 1e-8)
if err != nil {
t.Fatal(err)
}

if _, ok := ranks["c"]; ok {
t.Fatalf("expected 2-hop node c to be excluded, got ranks %v", ranks)
}
if _, ok := ranks["x"]; ok {
t.Fatalf("expected disconnected node x to be excluded, got ranks %v", ranks)
}
if _, ok := ranks["b"]; !ok {
t.Fatalf("expected 1-hop node b in local neighborhood, got ranks %v", ranks)
}
}

func TestPersonalizedPageRank_RespectsCanceledContext(t *testing.T) {
g := New()
addEdge(g, "a", "b", "r1", "knows", 1.0)

ctx, cancel := context.WithCancel(context.Background())
cancel()

_, err := g.PersonalizedPageRank(ctx, []string{"a"}, 2, 0.15, 20, 1e-6)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context.Canceled, got %v", err)
}
}
8 changes: 6 additions & 2 deletions internal/search/graph.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ func searchGraph(

switch opts.GraphScorer {
case GraphScorerPageRank:
return searchGraphPPR(ctx, store, g, kbID, seeds, seedSet, limit)
return searchGraphPPR(ctx, store, g, kbID, seeds, seedSet, opts, limit)
case GraphScorerWeighted:
return searchGraphWeighted(ctx, store, g, kbID, seeds, seedSet, opts, limit)
default:
Expand Down Expand Up @@ -110,9 +110,13 @@ func searchGraphPPR(
kbID string,
seeds []string,
seedSet map[string]struct{},
opts Options,
limit int,
) ([]*domain.SearchResult, error) {
ranks := g.PersonalizedPageRank(seeds, pprAlpha, pprMaxIter, pprEpsilon)
ranks, err := g.PersonalizedPageRank(ctx, seeds, opts.MaxHops, pprAlpha, pprMaxIter, pprEpsilon)
if err != nil {
return nil, fmt.Errorf("personalized pagerank: %w", err)
}

type entry struct {
id string
Expand Down
Loading