gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix graph query problem of PostgresSQL AGE

Browse Source
This commit is contained in:
yangdx
2025-04-25 16:04:41 +08:00
parent 36acc332b5
commit 5ca430adcd
1 changed files with 37 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
lightrag/kg/postgres_impl.py
View File

@@ -1879,11 +1879,14 @@ class PGGraphStorage(BaseGraphStorage):
visited_node_ids.add(internal_id) visited_node_ids.add(internal_id)
result.nodes.append(start_node) result.nodes.append(start_node)
while queue and len(visited_node_ids) < max_nodes: result.is_truncated = False
while queue:
# Dequeue the next node to process from the front of the queue # Dequeue the next node to process from the front of the queue
current_node, current_depth = queue.popleft() current_node, current_depth = queue.popleft()
if current_depth >= max_depth: # Check one more depth for backward edges
if current_depth > max_depth:
continue continue
# Get all edges and target nodes for the current node - query outgoing and incoming edges separately for efficiency # Get all edges and target nodes for the current node - query outgoing and incoming edges separately for efficiency
@@ -1935,6 +1938,9 @@ class PGGraphStorage(BaseGraphStorage):
properties=b_node["properties"], properties=b_node["properties"],
) )
# Sort entity_ids to ensure (A,B) and (B,A) are treated as the same edge
sorted_pair = tuple(sorted([current_entity_id, target_entity_id]))
# Create edge object # Create edge object
edge = KnowledgeGraphEdge( edge = KnowledgeGraphEdge(
id=edge_id, id=edge_id,
@@ -1944,38 +1950,38 @@ class PGGraphStorage(BaseGraphStorage):
properties=rel["properties"], properties=rel["properties"],
) )
# Sort entity_ids to ensure (A,B) and (B,A) are treated as the same edge if target_internal_id in visited_node_ids:
sorted_pair = tuple(sorted([current_entity_id, target_entity_id])) # Add backward edge if target node is visited
if (
edge_id not in visited_edges
and sorted_pair not in visited_edge_pairs
):
result.edges.append(edge)
visited_edges.add(edge_id)
visited_edge_pairs.add(sorted_pair)
# Add edge (if not already added) else:
if ( if len(visited_node_ids) < max_nodes and current_depth < max_depth:
edge_id not in visited_edges # If target node not yet visited, add to result and queue
and sorted_pair not in visited_edge_pairs result.nodes.append(target_node)
): visited_nodes.add(target_entity_id)
result.edges.append(edge) visited_node_ids.add(target_internal_id)
visited_edges.add(edge_id)
visited_edge_pairs.add(sorted_pair)
# If target node not yet visited, add to result and queue # Add node to queue with incremented depth
if target_internal_id not in visited_node_ids: queue.append((target_node, current_depth + 1))
result.nodes.append(target_node)
visited_nodes.add(target_entity_id)
visited_node_ids.add(target_internal_id)
# Add node to queue with incremented depth # Add forward edge
queue.append((target_node, current_depth + 1)) if (
edge_id not in visited_edges
# If node limit reached, set truncated flag and exit and sorted_pair not in visited_edge_pairs
if len(visited_node_ids) >= max_nodes: ):
result.is_truncated = True result.edges.append(edge)
logger.info( visited_edges.add(edge_id)
f"Graph truncated: BFS limited to {max_nodes} nodes" visited_edge_pairs.add(sorted_pair)
) # logger.info(f"Forward edge from {current_entity_id} to {target_entity_id}")
break else:
if current_depth < max_depth:
# If inner loop reached node limit and exited, also exit outer loop result.is_truncated = True
if len(visited_node_ids) >= max_nodes:
break
return result return result