Skip to content

Commit

Permalink
Add Gremlin group-by-depth (aws#251)
Browse files Browse the repository at this point in the history 
* Add --group-by-depth to Gremlin, modify OC with same explicit param

* Add tests

* PEP8 fixes

Co-authored-by: Michael Chin <[email protected]>
  • Loading branch information
@michaelnchin
michaelnchin and michaelnchin authored Feb 4, 2022
1 parent 76c3c04 commit 659b590
Show file tree
Hide file tree
Showing 6 changed files with 356 additions and 34 deletions.
6 changes: 6 additions & 0 deletions src/graph_notebook/magics/graph_magic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
parser.add_argument('-p', '--path-pattern', default='', help='path pattern')
parser.add_argument('-g', '--group-by', type=str, default='T.label',
help='Property used to group nodes (e.g. code, T.region) default is T.label')
parser.add_argument('-gd', '--group-by-depth', action='store_true', default=False,
help="Group nodes based on path hierarchy")
parser.add_argument('-gr', '--group-by-raw', action='store_true', default=False,
help="Group nodes by the raw result")
parser.add_argument('-d', '--display-property', type=str, default='T.label',
Expand Down Expand Up @@ -578,6 +580,7 @@ def gremlin(self, line, cell, local_ns: dict = None):
logger.debug(f'ignore_groups: {args.ignore_groups}')
gn = GremlinNetwork(group_by_property=args.group_by, display_property=args.display_property,
group_by_raw=args.group_by_raw,
group_by_depth=args.group_by_depth,
edge_display_property=args.edge_display_property,
tooltip_property=args.tooltip_property,
edge_tooltip_property=args.edge_tooltip_property,
Expand Down Expand Up @@ -1665,6 +1668,8 @@ def handle_opencypher_query(self, line, cell, local_ns):
parser = argparse.ArgumentParser()
parser.add_argument('-g', '--group-by', type=str, default='~labels',
help='Property used to group nodes (e.g. code, ~id) default is ~labels')
parser.add_argument('-gd', '--group-by-depth', action='store_true', default=False,
help="Group nodes based on path hierarchy")
parser.add_argument('-gr', '--group-by-raw', action='store_true', default=False,
help="Group nodes by the raw result")
parser.add_argument('mode', nargs='?', default='query', help='query mode [query|bolt]',
Expand Down Expand Up @@ -1721,6 +1726,7 @@ def handle_opencypher_query(self, line, cell, local_ns):
try:
gn = OCNetwork(group_by_property=args.group_by, display_property=args.display_property,
group_by_raw=args.group_by_raw,
group_by_depth=args.group_by_depth,
edge_display_property=args.edge_display_property,
tooltip_property=args.tooltip_property,
edge_tooltip_property=args.edge_tooltip_property,
Expand Down
1 change: 1 addition & 0 deletions src/graph_notebook/network/EventfulNetwork.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
DEFAULT_GRP = 'DEFAULT_GROUP'
DEFAULT_RAW_GRP_KEY = '__RAW_RESULT__'
DEFAULT_LABEL_MAX_LENGTH = 10
DEPTH_GRP_KEY = 'TRAVERSAL_DEPTH'

VALID_EVENTS = [EVENT_ADD_NODE, EVENT_ADD_NODE_DATA, EVENT_ADD_NODE_PROPERTY, EVENT_ADD_EDGE, EVENT_ADD_EDGE_DATA]

Expand Down
54 changes: 35 additions & 19 deletions src/graph_notebook/network/gremlin/GremlinNetwork.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import logging
from enum import Enum

from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP, DEFAULT_RAW_GRP_KEY
from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP, DEPTH_GRP_KEY, DEFAULT_RAW_GRP_KEY
from gremlin_python.process.traversal import T, Direction
from gremlin_python.structure.graph import Path, Vertex, Edge
from networkx import MultiDiGraph
Expand Down Expand Up @@ -100,9 +100,11 @@ class GremlinNetwork(EventfulNetwork):
def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=DEFAULT_LABEL_MAX_LENGTH,
edge_label_max_length=DEFAULT_LABEL_MAX_LENGTH, group_by_property=T_LABEL, display_property=T_LABEL,
edge_display_property=T_LABEL, tooltip_property=None, edge_tooltip_property=None, ignore_groups=False,
group_by_raw=False):
group_by_depth=False, group_by_raw=False):
if graph is None:
graph = MultiDiGraph()
if group_by_depth:
group_by_property = DEPTH_GRP_KEY
super().__init__(graph, callbacks, label_max_length, edge_label_max_length, group_by_property,
display_property, edge_display_property, tooltip_property, edge_tooltip_property,
ignore_groups, group_by_raw)
Expand Down Expand Up @@ -220,10 +222,9 @@ def add_results_with_pattern(self, results, pattern_list: list):
raise INVALID_VERTEX_PATH_PATTERN_ERROR

# add the vertex, no matter what patterns border a vertex pattern, it will always be added.
self.add_vertex(path[path_index])
self.add_vertex(path[path_index], path_index=i)
# if the path index is over 0, we need to handle edges between this node and the previous one
if path_index > 0:
# TODO: First node (PathPattern.V) is rendered blank if T.id is a integer instead of a string.
if path_pattern == PathPattern.V:
# two V patterns next to each other is an undirected, unlabeled edge
if previous_pattern == PathPattern.V:
Expand Down Expand Up @@ -294,7 +295,7 @@ def add_results(self, results):
if not isinstance(results, list):
raise ValueError("results must be a list of paths")

for path in results:
for path_index, path in enumerate(results):
if isinstance(path, Path):
if type(path[0]) is Edge or type(path[-1]) is Edge:
raise INVALID_PATH_ERROR
Expand Down Expand Up @@ -324,18 +325,20 @@ def add_results(self, results):
else:
self.insert_path_element(path, i)
elif isinstance(path, dict) and T.id in path.keys() and T.label in path.keys():
self.insert_elementmap(path)
self.insert_elementmap(path, index=path_index)
else:
raise ValueError("all entries in results must be paths or elementMaps")

def add_vertex(self, v):
def add_vertex(self, v, path_index: int = -1):
"""
Adds a vertex to the network. If v is of :type Vertex, we will gather its id and label.
If v comes from a valueMap, it will be a dict, with the keys T.label and T.ID being present in
the dict for gathering the label and id, respectively.
:param v: The vertex taken from a path traversal object.
:param path_index: Position of the element in the path traversal order
"""
depth_group = "__DEPTH-" + str(path_index) + "__"
node_id = ''
if type(v) is Vertex:
node_id = v.id
Expand All @@ -357,6 +360,8 @@ def add_vertex(self, v):
group = v.label
elif str(self.group_by_property) in [T_ID, 'id']:
group = v.id
elif self.group_by_property == DEPTH_GRP_KEY:
group = depth_group
else:
group = DEFAULT_GRP
else: # handle dict format group_by
Expand All @@ -368,6 +373,8 @@ def add_vertex(self, v):
group = v.label
elif self.group_by_property[str(v.label)] in [T_ID, 'id']:
group = v.id
elif self.group_by_property[str(v.label)] == DEPTH_GRP_KEY:
group = depth_group
else:
group = vertex_dict[self.group_by_property[str(v.label)]]
else:
Expand All @@ -379,7 +386,7 @@ def add_vertex(self, v):

if self.display_property:
label_property_raw_value = self.get_explicit_vertex_property_value(node_id, title,
self.display_property)
self.display_property)
if label_property_raw_value:
label_full, label = self.strip_and_truncate_label_and_title(label_property_raw_value,
self.label_max_length)
Expand Down Expand Up @@ -432,17 +439,24 @@ def add_vertex(self, v):
if not group_is_set:
if isinstance(self.group_by_property, dict):
try:
if self.group_by_property[title_plc] == DEFAULT_RAW_GRP_KEY:
if DEPTH_GRP_KEY == self.group_by_property[title_plc]:
group = depth_group
group_is_set = True
elif DEFAULT_RAW_GRP_KEY == self.group_by_property[title_plc]:
group = str(v)
group_is_set = True
elif str(k) == self.group_by_property[title_plc]:
group = str(v[k])
group_is_set = True
except KeyError:
pass
elif str(k) == self.group_by_property:
group = str(v[k])
group_is_set = True
else:
if DEPTH_GRP_KEY == self.group_by_property:
group = depth_group
group_is_set = True
elif str(k) == self.group_by_property:
group = str(v[k])
group_is_set = True
if not display_is_set:
label_property_raw_value = self.get_dict_element_property_value(v, k, title_plc,
self.display_property)
Expand Down Expand Up @@ -480,7 +494,9 @@ def add_vertex(self, v):
else:
node_id = str(v)
title = str(v)
if str(self.group_by_property) == DEFAULT_RAW_GRP_KEY:
if self.group_by_property == DEPTH_GRP_KEY:
group = depth_group
if self.group_by_property == DEFAULT_RAW_GRP_KEY:
group = str(v)
else:
group = DEFAULT_GRP
Expand Down Expand Up @@ -605,7 +621,7 @@ def add_blank_edge(self, from_id, to_id, edge_id=None, undirected=True, label=''

def insert_path_element(self, path, i):
if i == 0:
self.add_vertex(path[i])
self.add_vertex(path[i], path_index=0)
return

if type(path[i]) is Edge:
Expand Down Expand Up @@ -634,7 +650,7 @@ def insert_path_element(self, path, i):
else:
from_id = get_id(path[i - 1])

self.add_vertex(path[i])
self.add_vertex(path[i], path_index=i)
if type(path[i - 1]) is not Edge:
if type(path[i - 1]) is dict:
if Direction.IN not in path[i-1]:
Expand All @@ -658,14 +674,14 @@ def insert_elementmap(self, e_map, check_emap=False, path_element=None, index=No
# Ensure that the default nodes includes with edge elementMaps do not overwrite nodes
# with the same ID that have been explicitly inserted.
if not self.graph.has_node(from_id):
self.add_vertex(e_map[Direction.OUT])
self.add_vertex(e_map[Direction.OUT], path_index=index-1)
if not self.graph.has_node(to_id):
self.add_vertex(e_map[Direction.IN])
self.add_vertex(e_map[Direction.IN], path_index=index+1)
self.add_path_edge(e_map, from_id, to_id)
# Handle vertex elementMap
else:
# Overwrite the the default node created by edge elementMap, if it exists already.
# If a default node was created at the same index by an edge elementMap, overwrite it.
if check_emap:
self.insert_path_element(path_element, index)
else:
self.add_vertex(e_map)
self.add_vertex(e_map, path_index=index)
20 changes: 11 additions & 9 deletions src/graph_notebook/network/opencypher/OCNetwork.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,9 @@
SPDX-License-Identifier: Apache-2.0
"""

import json
import logging

from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP, DEFAULT_RAW_GRP_KEY
from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP, DEPTH_GRP_KEY, DEFAULT_RAW_GRP_KEY
from networkx import MultiDiGraph

logging.basicConfig()
Expand Down Expand Up @@ -35,9 +34,12 @@ def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=
edge_label_max_length=DEFAULT_LABEL_MAX_LENGTH, group_by_property=LABEL_KEY,
display_property=LABEL_KEY, edge_display_property=EDGE_TYPE_KEY,
tooltip_property=None, edge_tooltip_property=None,
ignore_groups=False, group_by_raw=False):
ignore_groups=False,
group_by_depth=False, group_by_raw=False):
if graph is None:
graph = MultiDiGraph()
if group_by_depth:
group_by_property = DEPTH_GRP_KEY
super().__init__(graph, callbacks, label_max_length, edge_label_max_length, group_by_property,
display_property, edge_display_property, tooltip_property, edge_tooltip_property,
ignore_groups, group_by_raw)
Expand Down Expand Up @@ -127,14 +129,14 @@ def parse_node(self, node: dict, path_index: int = -2):

if not isinstance(self.group_by_property, dict): # Handle string format group_by
try:
if self.group_by_property == DEFAULT_RAW_GRP_KEY:
if self.group_by_property == DEPTH_GRP_KEY:
group = depth_group
elif self.group_by_property == DEFAULT_RAW_GRP_KEY:
group = str(node)
elif self.group_by_property in [LABEL_KEY, 'labels'] and len(node[LABEL_KEY]) > 0:
group = node[LABEL_KEY][0]
elif self.group_by_property in [ID_KEY, 'id']:
group = node[ID_KEY]
elif self.group_by_property == "TRAVERSAL_DEPTH":
group = depth_group
elif self.group_by_property in node[PROPERTIES_KEY]:
group = node[PROPERTIES_KEY][self.group_by_property]
else:
Expand All @@ -145,14 +147,14 @@ def parse_node(self, node: dict, path_index: int = -2):
try:
if str(node[LABEL_KEY][0]) in self.group_by_property and len(node[LABEL_KEY]) > 0:
key = node[LABEL_KEY][0]
if self.group_by_property[key] == DEFAULT_RAW_GRP_KEY:
if self.group_by_property[key] == DEPTH_GRP_KEY:
group = depth_group
elif self.group_by_property[key] == DEFAULT_RAW_GRP_KEY:
group = str(node)
elif self.group_by_property[key] in [LABEL_KEY, 'labels']:
group = node[LABEL_KEY][0]
elif self.group_by_property[key] in [ID_KEY, 'id']:
group = node[ID_KEY]
elif self.group_by_property[key] == "TRAVERSAL_DEPTH":
group = depth_group
else:
group = node[PROPERTIES_KEY][self.group_by_property[key]]
else:
Expand Down
Loading

0 comments on commit 659b590

Please sign in to comment.