diff --git a/nebula3/data/ResultSet.py b/nebula3/data/ResultSet.py index b9a5954b..2f87a8d9 100644 --- a/nebula3/data/ResultSet.py +++ b/nebula3/data/ResultSet.py @@ -8,7 +8,7 @@ from nebula3.common.ttypes import ErrorCode -from nebula3.data.DataObject import DataSetWrapper +from nebula3.data.DataObject import DataSetWrapper, Node, Relationship, PathWrapper class ResultSet(object): @@ -191,6 +191,170 @@ def rows(self): if self._data_set_wrapper is None: return [] return self._data_set_wrapper.get_rows() + + def dict_for_vis(self): + """Convert result set to a dictionary format suitable for visualization. + + Example: + { + "nodes":[ + { + "id":"player100", + "labels":[ + "player" + ], + "props":{ + "name":"Tim Duncan", + "age":"42", + "id":"player100" + } + }, + { + "id":"player101", + "labels":[ + "player" + ], + "props":{ + "age":"36", + "name":"Tony Parker", + "id":"player101" + } + } + ], + "edges":[ + { + "src":"player100", + "dst":"player101", + "name":"follow", + "rank":0, + "props":{ + "degree":"95" + } + } + ], + "nodes_dict":{ + "player100":{ + "id":"player100", + "labels":[ + "player" + ], + "props":{ + "name":"Tim Duncan", + "age":"42", + "id":"player100" + } + }, + "player101":{ + "id":"player101", + "labels":[ + "player" + ], + "props":{ + "age":"36", + "name":"Tony Parker", + "id":"player101" + } + } + }, + "edges_dict":{ + "(""player100", + "player101", + 0, + "follow"")":{ + "src":"player100", + "dst":"player101", + "name":"follow", + "rank":0, + "props":{ + "degree":"95" + } + } + }, + "nodes_count":2, + "edges_count":1 + } + + :return: dict with keys: + nodes, edges, nodes_dict, edges_dict, nodes_count, edges_count + """ + def add_to_nodes_or_edges(nodes_dict, edges_dict, item): + if isinstance(item, Node): + node_id = str(item.get_id().cast()) + tags = item.tags() # list of strings + props_raw = dict() + for tag in tags: + # TODO: handle duplicate keys + props_raw.update(item.properties(tag)) + props = { + k: str(v.cast()) if hasattr(v, "cast") else str(v) + for k, v in props_raw.items() + } + + if "id" not in props: + props["id"] = node_id + + if node_id not in nodes_dict: + nodes_dict[node_id] = { + "id": node_id, + "labels": tags, + "props": props, + } + else: + nodes_dict[node_id]["labels"] = list(set(nodes_dict[node_id]["labels"] + tags)) + nodes_dict[node_id]["props"].update(props) + + elif isinstance(item, Relationship): + src_id = str(item.start_vertex_id().cast()) + dst_id = str(item.end_vertex_id().cast()) + rank = item.ranking() + edge_name = item.edge_name() + props_raw = item.properties() + props = { + k: str(v.cast()) if hasattr(v, "cast") else str(v) + for k, v in props_raw.items() + } + if (src_id, dst_id, rank, edge_name) not in edges_dict: + edges_dict[(src_id, dst_id, rank, edge_name)] = { + "src": src_id, + "dst": dst_id, + "name": edge_name, + "rank": rank, + "props": props, + } + else: + edges_dict[(src_id, dst_id, rank, edge_name)]["props"].update(props) + + elif isinstance(item, PathWrapper): + for node in item.nodes(): + add_to_nodes_or_edges(nodes_dict, edges_dict, node) + for edge in item.relationships(): + add_to_nodes_or_edges(nodes_dict, edges_dict, edge) + + elif isinstance(item, list): + for it in item: + add_to_nodes_or_edges(nodes_dict, edges_dict, it) + + nodes_dict = dict() + edges_dict = dict() + + columns = self.keys() + for col_num in range(self.col_size()): + col_name = columns[col_num] + col_list = self.column_values(col_name) + add_to_nodes_or_edges( + nodes_dict, edges_dict, [x.cast() for x in col_list] + ) + nodes = list(nodes_dict.values()) + edges = list(edges_dict.values()) + + return { + "nodes": nodes, + "edges": edges, + "nodes_dict": nodes_dict, + "edges_dict": edges_dict, + "nodes_count": len(nodes), + "edges_count": len(edges), + } def __iter__(self): """the iterator for per row