Revert "refactor: use get_nerghbors in storage instead of building adj_list"

This reverts commit e924945.

Author: ChenZiHong-Gavin

graphgen/bases/base_partitioner.py

@@ -51,3 +51,24 @@ def community2batch(
             if edge_data:
                 edges_data.append((u, v, edge_data))
         return nodes_data, edges_data
+
+    @staticmethod
+    def _build_adjacency_list(
+        nodes: List[tuple[str, dict]], edges: List[tuple[str, str, dict]]
+    ) -> tuple[dict[str, List[str]], set[tuple[str, str]]]:
+        """
+        Build adjacency list and edge set from nodes and edges.
+        :param nodes
+        :param edges
+        :return: adjacency list, edge set
+        """
+        adj: dict[str, List[str]] = {n[0]: [] for n in nodes}
+        edge_set: set[tuple[str, str]] = set()
+        for u, v, _ in edges:
+            if u == v:
+                continue
+            adj[u].append(v)
+            adj[v].append(u)
+            edge_set.add((u, v))
+            edge_set.add((v, u))
+        return adj, edge_set

graphgen/bases/base_storage.py

@@ -126,10 +126,6 @@ def upsert_edge(
     def delete_node(self, node_id: str):
         raise NotImplementedError
 
-    @abstractmethod
-    def get_neighbors(self, node_id: str) -> List[str]:
-        raise NotImplementedError
-
     @abstractmethod
     def reload(self):
         raise NotImplementedError

graphgen/common/init_storage.py

@@ -129,9 +129,6 @@ def upsert_edge(
     def delete_node(self, node_id: str):
         return self.graph.delete_node(node_id)
 
-    def get_neighbors(self, node_id: str) -> List[str]:
-        return self.graph.get_neighbors(node_id)
-
     def reload(self):
         return self.graph.reload()
 
@@ -248,9 +245,6 @@ def upsert_edge(
     def delete_node(self, node_id: str):
         return ray.get(self.actor.delete_node.remote(node_id))
 
-    def get_neighbors(self, node_id: str) -> List[str]:
-        return ray.get(self.actor.get_neighbors.remote(node_id))
-
     def reload(self):
         return ray.get(self.actor.reload.remote())
 

graphgen/models/partitioner/bfs_partitioner.py

@@ -26,6 +26,8 @@ def partition(
         nodes = g.get_all_nodes()
         edges = g.get_all_edges()
 
+        adj, _ = self._build_adjacency_list(nodes, edges)
+
         used_n: set[str] = set()
         used_e: set[frozenset[str]] = set()
 
@@ -53,16 +55,15 @@ def partition(
                     used_n.add(it)
                     comm_n.append(it)
                     cnt += 1
-                    for nei in g.get_neighbors(it):
+                    for nei in adj[it]:
                         e_key = frozenset((it, nei))
                         if e_key not in used_e:
                             queue.append((EDGE_UNIT, e_key))
                 else:
                     if it in used_e:
                         continue
                     used_e.add(it)
-                    u, v = tuple(it)
-                    comm_e.append((u, v))
+                    comm_e.append(tuple(sorted(it)))
                     cnt += 1
                     # push nodes that are not visited
                     for n in it:

graphgen/models/partitioner/dfs_partitioner.py

@@ -26,6 +26,8 @@ def partition(
         nodes = g.get_all_nodes()
         edges = g.get_all_edges()
 
+        adj, _ = self._build_adjacency_list(nodes, edges)
+
         used_n: set[str] = set()
         used_e: set[frozenset[str]] = set()
 
@@ -53,7 +55,7 @@ def partition(
                     used_n.add(it)
                     comm_n.append(it)
                     cnt += 1
-                    for nei in g.get_neighbors(it):
+                    for nei in adj[it]:
                         e_key = frozenset((it, nei))
                         if e_key not in used_e:
                             stack.append((EDGE_UNIT, e_key))
@@ -62,8 +64,7 @@ def partition(
                     if it in used_e:
                         continue
                     used_e.add(it)
-                    u, v = sorted(it)
-                    comm_e.append((u, v))
+                    comm_e.append(tuple(sorted(it)))
                     cnt += 1
                     # push neighboring nodes
                     for n in it:

graphgen/models/partitioner/ece_partitioner.py

@@ -65,6 +65,7 @@ def partition(
         nodes: List[Tuple[str, dict]] = g.get_all_nodes()
         edges: List[Tuple[str, str, dict]] = g.get_all_edges()
 
+        adj, _ = self._build_adjacency_list(nodes, edges)
         node_dict = dict(nodes)
         edge_dict = {frozenset((u, v)): d for u, v, d in edges}
 
@@ -117,7 +118,7 @@ def _add_unit(u):
 
                 neighbors: List[Tuple[str, Any, dict]] = []
                 if cur_type == NODE_UNIT:
-                    for nb_id in g.get_neighbors(cur_id):
+                    for nb_id in adj.get(cur_id, []):
                         e_key = frozenset((cur_id, nb_id))
                         if e_key not in used_e and e_key not in community_edges:
                             neighbors.append((EDGE_UNIT, e_key, edge_dict[e_key]))

graphgen/models/storage/graph/kuzu_storage.py

@@ -373,19 +373,6 @@ def delete_node(self, node_id: str):
         self._conn.execute(query, {"id": node_id})
         print(f"Node {node_id} deleted from KuzuDB.")
 
-    def get_neighbors(self, node_id: str) -> List[str]:
-        query = """
-            MATCH (a:Entity {id: $id})-[:Relation]->(b:Entity)
-            RETURN b.id
-        """
-        result = self._conn.execute(query, {"id": node_id})
-        neighbors = []
-        while result.has_next():
-            row = result.get_next()
-            if row and len(row) >= 1:
-                neighbors.append(row[0])
-        return neighbors
-
     def clear(self):
         """Clear all data but keep schema (or drop tables)."""
         self._conn.execute("MATCH (n) DETACH DELETE n")

graphgen/models/storage/graph/networkx_storage.py

@@ -188,18 +188,6 @@ def delete_node(self, node_id: str):
         else:
             print(f"Node {node_id} not found in the graph for deletion.")
 
-    def get_neighbors(self, node_id: str) -> List[str]:
-        """
-        Get the neighbors of a node based on the specified node_id.
-
-        :param node_id: The node_id to get neighbors for
-        :return: List of neighbor node IDs
-        """
-        if self._graph.has_node(node_id):
-            return list(self._graph.neighbors(node_id))
-        print(f"Node {node_id} not found in the graph.")
-        return []
-
     def clear(self):
         """
         Clear the graph by removing all nodes and edges.