Add basic node property support

Author: knutwalker

graphdatascience/graph/graph_cypher_runner.py

@@ -1,4 +1,4 @@
-from collections import namedtuple
+from collections import defaultdict, namedtuple
 from typing import Any, NamedTuple, Optional, Tuple
 
 from pandas import Series
@@ -62,6 +62,12 @@ def __str__(self) -> str:
         return f"{self.left_arrow}{self.type_filter}{self.right_arrow}(target{self.label_filter})"
 
 
+class LabelPropertyMapping(NamedTuple):
+    label: str
+    property_key: str
+    default_value: Optional[Any] = None
+
+
 class GraphCypherRunner(IllegalAttrChecker):
     def __init__(self, query_runner: QueryRunner, namespace: str, server_version: ServerVersion) -> None:
         if server_version < ServerVersion(2, 4, 0):
@@ -121,6 +127,8 @@ def project(
             right_arrow="-" if inverse else "->",
         )
 
+        label_mappings = defaultdict(list)
+
         if nodes:
             if len(nodes) == 1 or combine_labels_with == "AND":
                 match_pattern = match_pattern._replace(label_filter=f":{':'.join(spec.source_label for spec in nodes)}")
@@ -147,14 +155,22 @@ def project(
             else:
                 raise ValueError(f"Invalid value for combine_labels_with: {combine_labels_with}")
 
+            for spec in nodes:
+                if spec.properties:
+                    for prop in spec.properties:
+                        label_mappings[spec.source_label].append(
+                            LabelPropertyMapping(spec.source_label, prop.property_key, prop.default_value)
+                        )
+
+        rel_var = ""
         if rels:
             if len(rels) == 1:
-                rel_var = ""
                 data_config["relationshipType"] = rels[0].source_type
             else:
                 rel_var = "rel"
                 data_config["relationshipTypes"] = "type(rel)"
                 data_config_is_static = False
+
             match_pattern = match_pattern._replace(
                 type_filter=f"[{rel_var}:{'|'.join(spec.source_type for spec in rels)}]"
             )
@@ -169,6 +185,24 @@ def project(
 
         match_part = str(match_part)
 
+        case_part = []
+        if label_mappings:
+            with_rel = f", {rel_var}" if rel_var else ""
+            case_part = [f"WITH source, target{with_rel}"]
+            for kind in ["source", "target"]:
+                case_part.append("CASE")
+
+                for label, mappings in label_mappings.items():
+                    mappings = ", ".join(f".{key.property_key}" for key in mappings)
+                    when_part = f"WHEN '{label}' in labels({kind}) THEN [{kind} {{{mappings}}}]"
+                    case_part.append(when_part)
+
+                case_part.append(f"END AS {kind}NodeProperties")
+
+            data_config["sourceNodeProperties"] = "sourceNodeProperties"
+            data_config["targetNodeProperties"] = "targetNodeProperties"
+            data_config_is_static = False
+
         args = ["$graph_name", "source", "target"]
 
         if data_config:
@@ -184,9 +218,7 @@ def project(
 
         return_part = f"RETURN {self._namespace}({', '.join(args)})"
 
-        query = "\n".join(part for part in [match_part, return_part] if part)
-
-        print(query)
+        query = "\n".join(part for part in [match_part, *case_part, return_part] if part)
 
         result = self._query_runner.run_query_with_logging(
             query,
@@ -208,16 +240,39 @@ def _node_projections_spec(self, spec: Any) -> list[NodeProjection]:
         if isinstance(spec, dict):
             return [self._node_projection_spec(node, name) for name, node in spec.items()]
 
-        raise TypeError(f"Invalid node projection specification: {spec}")
+        raise TypeError(f"Invalid node projections specification: {spec}")
 
     def _node_projection_spec(self, spec: Any, name: Optional[str] = None) -> NodeProjection:
         if isinstance(spec, str):
             return NodeProjection(name=name or spec, source_label=spec)
 
+        if name is None:
+            raise ValueError(f"Node projections with properties must use the dict syntax: {spec}")
+
+        if isinstance(spec, dict):
+            properties = [self._node_properties_spec(prop, name) for name, prop in spec.items()]
+            return NodeProjection(name=name, source_label=name, properties=properties)
+
+        if isinstance(spec, list):
+            properties = [self._node_properties_spec(prop) for prop in spec]
+            return NodeProjection(name=name, source_label=name, properties=properties)
+
         raise TypeError(f"Invalid node projection specification: {spec}")
 
-    def _node_properties_spec(self, properties: dict[str, Any]) -> list[NodeProperty]:
-        raise TypeError(f"Invalid node projection specification: {properties}")
+    def _node_properties_spec(self, spec: Any, name: Optional[str] = None) -> NodeProperty:
+        if isinstance(spec, str):
+            return NodeProperty(name=name or spec, property_key=spec)
+
+        if name is None:
+            raise ValueError(f"Node properties spec must be used with the dict syntax: {spec}")
+
+        if spec is True:
+            return NodeProperty(name=name, property_key=name)
+
+        if isinstance(spec, dict):
+            return NodeProperty(name=name, property_key=name, **spec)
+
+        raise TypeError(f"Invalid node property specification: {spec}")
 
     def _rel_projections_spec(self, spec: Any) -> list[RelationshipProjection]:
         if spec is None or spec is False:

graphdatascience/tests/unit/test_graph_cypher.py

@@ -216,3 +216,34 @@ def test_multiple_multi_graph(runner: CollectingQueryRunner, gds: GraphDataScien
         "targetNodeLabels: labels(target), "
         "relationshipTypes: type(rel)})"
     )
+
+
+@pytest.mark.parametrize("server_version", [ServerVersion(2, 4, 0)])
+def test_node_properties(runner: CollectingQueryRunner, gds: GraphDataScience) -> None:
+    G, _ = gds.graph.cypher.project(
+        "g", nodes=dict(L1=["prop1"], L2=["prop2", "prop3"], L3=dict(prop4=True, prop5=dict()))
+    )
+
+    assert G.name() == "g"
+    assert runner.last_params() == dict(graph_name="g")
+
+    assert runner.last_query() == (
+        """MATCH (source)-->(target)
+WHERE (source:L1 OR source:L2 OR source:L3) AND (target:L1 OR target:L2 OR target:L3)
+WITH source, target
+CASE
+WHEN 'L1' in labels(source) THEN [source {.prop1}]
+WHEN 'L2' in labels(source) THEN [source {.prop2, .prop3}]
+WHEN 'L3' in labels(source) THEN [source {.prop4, .prop5}]
+END AS sourceNodeProperties
+CASE
+WHEN 'L1' in labels(target) THEN [target {.prop1}]
+WHEN 'L2' in labels(target) THEN [target {.prop2, .prop3}]
+WHEN 'L3' in labels(target) THEN [target {.prop4, .prop5}]
+END AS targetNodeProperties
+RETURN gds.graph.project($graph_name, source, target, {"""
+        "sourceNodeLabels: labels(source), "
+        "targetNodeLabels: labels(target), "
+        "sourceNodeProperties: sourceNodeProperties, "
+        "targetNodeProperties: targetNodeProperties})"
+    )