From ea663a6ef3ae91b2b416fd3d23bea6e9bbff16d9 Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Mon, 6 Apr 2026 23:28:52 +0700 Subject: [PATCH 1/2] Localize PersonalizedPageRank traversal --- internal/graph/graph.go | 176 +++++++++++++++++++++++++---------- internal/graph/graph_test.go | 53 ++++++++++- internal/search/graph.go | 8 +- 3 files changed, 185 insertions(+), 52 deletions(-) diff --git a/internal/graph/graph.go b/internal/graph/graph.go index 55b1591..d3c49a2 100644 --- a/internal/graph/graph.go +++ b/internal/graph/graph.go @@ -1,6 +1,7 @@ package graph import ( + "context" "math" "sync" "time" @@ -178,8 +179,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. @@ -389,70 +390,89 @@ 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 } - // 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 } - // 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 _, nextID := range neighbors[id] { + 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 { @@ -466,7 +486,69 @@ 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 _, e := range edges { + nodes[e.TargetID] = struct{}{} + if _, seen := dist[e.TargetID]; seen { + continue + } + 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) + } + for _, id := range nodeIDs { + localNeighbors := make([]string, 0, len(g.forward[id])+len(g.reverse[id])) + for _, edges := range [][]Edge{g.forward[id], g.reverse[id]} { + for _, e := range edges { + if _, ok := nodes[e.TargetID]; !ok { + continue + } + localNeighbors = append(localNeighbors, e.TargetID) + } + } + neighbors[id] = localNeighbors + } + + return nodeIDs, neighbors, nil } // --- Helpers --- diff --git a/internal/graph/graph_test.go b/internal/graph/graph_test.go index 1311446..31b481f 100644 --- a/internal/graph/graph_test.go +++ b/internal/graph/graph_test.go @@ -1,6 +1,8 @@ package graph import ( + "context" + "errors" "testing" "time" ) @@ -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"]) @@ -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"]) @@ -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) + } +} diff --git a/internal/search/graph.go b/internal/search/graph.go index b6c971d..427e645 100644 --- a/internal/search/graph.go +++ b/internal/search/graph.go @@ -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: @@ -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 From e392a168b59a2ab4bf765e6391aa0ae4ce34b876 Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Tue, 7 Apr 2026 14:30:28 +0700 Subject: [PATCH 2/2] Address PR review feedback --- internal/graph/graph.go | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/internal/graph/graph.go b/internal/graph/graph.go index d3c49a2..13d5865 100644 --- a/internal/graph/graph.go +++ b/internal/graph/graph.go @@ -3,6 +3,7 @@ package graph import ( "context" "math" + "sort" "sync" "time" ) @@ -461,7 +462,12 @@ func (g *Graph) PersonalizedPageRank( continue } share := rank[id] / float64(len(neighbors[id])) - for _, nextID := range neighbors[id] { + for j, nextID := range neighbors[id] { + if j%256 == 0 { + if err := ctx.Err(); err != nil { + return nil, err + } + } newRank[nextID] += share } } @@ -519,7 +525,12 @@ func (g *Graph) localPPRNeighborhoodLocked( continue } for _, edges := range [][]Edge{g.forward[cur], g.reverse[cur]} { - for _, e := range edges { + 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 @@ -535,10 +546,23 @@ func (g *Graph) localPPRNeighborhoodLocked( for id := range nodes { nodeIDs = append(nodeIDs, id) } - for _, id := range nodeIDs { + 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 }