Skip to content

fix: add special edge handling in bfs, dfs and base partitioner #150

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

Merged
ChenZiHong-Gavin merged 5 commits into from
Jan 8, 2026
Merged

fix: add special edge handling in bfs, dfs and base partitioner #150

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eea3e2a
fix: add special edge handling in bfs, dfs and base partitioner
CHERRY-ui8 Jan 5, 2026
c534f94
refactor: extract a helper method in the base class
CHERRY-ui8 Jan 5, 2026
4b55d10
fix: add List import to dfs_partitioner
CHERRY-ui8 Jan 5, 2026
33e0afb
fix: align edge type
ChenZiHong-Gavin Jan 8, 2026
b7c675c
fix: fix lint error
ChenZiHong-Gavin Jan 8, 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
7 changes: 6 additions & 1 deletion graphgen/bases/base_partitioner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,12 @@ def community2batch(
if node_data:
nodes_data.append((node, node_data))
edges_data = []
for u, v in edges:
for edge in edges:
# Filter out self-loops and invalid edges
if isinstance(edge, (tuple, list)) and len(edge) == 2:
u, v = edge[0], edge[1]
else:
continue
edge_data = g.get_edge(u, v)
if edge_data:
edges_data.append((u, v, edge_data))
Expand Down
13 changes: 8 additions & 5 deletions graphgen/models/partitioner/bfs_partitioner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def partition(
continue

comm_n: List[str] = []
comm_e: List[tuple[str, str]] = []
comm_e: List[frozenset[str]] = []
queue: deque[tuple[str, Any]] = deque([(kind, seed)])
cnt = 0

Expand All @@ -63,14 +63,17 @@ def partition(
if it in used_e:
continue
used_e.add(it)

u, v = it
comm_e.append((u, v))
comm_e.append(it)
cnt += 1
# push nodes that are not visited
for n in it:
if n not in used_n:
queue.append((NODE_UNIT, n))

if comm_n or comm_e:
yield Community(id=seed, nodes=comm_n, edges=comm_e)
# Filter out self-loops and invalid edges
valid_edges = [
tuple(edge) for edge in comm_e
if isinstance(edge, frozenset) and len(edge) == 2
]
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jan 5, 2026

Choose a reason for hiding this comment

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

medium

This logic for filtering out invalid edges is duplicated in graphgen/models/partitioner/dfs_partitioner.py (lines 76-79). To improve maintainability and adhere to the DRY (Don't Repeat Yourself) principle, consider extracting this list comprehension into a common helper function. Since both BFSPartitioner and DFSPartitioner inherit from BasePartitioner, a static or protected helper method in the base class would be a suitable location.

yield Community(id=seed, nodes=comm_n, edges=valid_edges)
12 changes: 9 additions & 3 deletions graphgen/models/partitioner/dfs_partitioner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,8 @@ def partition(
):
continue

comm_n, comm_e = [], []
comm_n: List[str] = []
comm_e: List[frozenset[str]] = []
stack = [(kind, seed)]
cnt = 0

Expand All @@ -63,12 +64,17 @@ def partition(
if it in used_e:
continue
used_e.add(it)
comm_e.append(tuple(it))
comm_e.append(it)
cnt += 1
# push neighboring nodes
for n in it:
if n not in used_n:
stack.append((NODE_UNIT, n))

if comm_n or comm_e:
yield Community(id=seed, nodes=comm_n, edges=comm_e)
# Filter out self-loops and invalid edges
valid_edges = [
tuple(edge) for edge in comm_e
if isinstance(edge, frozenset) and len(edge) == 2
]
yield Community(id=seed, nodes=comm_n, edges=valid_edges)
Loading