Merge pull request #122 from adamnsch/allow-string-label-in-construct

Allow string label in community `construct`

Author: FlorentinD

README.md

@@ -89,17 +89,21 @@ To do this we provide the `gds.alpha.graph.construct` method with node data fram
 
 ```python
 nodes = pandas.DataFrame(
-    "nodeId": [0, 1, 2, 3],
-    "labels":  ["A", "B", "C", "A"],
-    "prop1": [42, 1337, 8, 0],
-    "otherProperty": [0.1, 0.2, 0.3, 0.4]
+    {
+        "nodeId": [0, 1, 2, 3],
+        "labels":  ["A", "B", "C", "A"],
+        "prop1": [42, 1337, 8, 0],
+        "otherProperty": [0.1, 0.2, 0.3, 0.4]
+    }
 )
 
 relationships = pandas.DataFrame(
-    "sourceId": [0, 1, 2, 3],
-    "targetId": [1, 2, 3, 0],
-    "relationshipType": ["REL", "REL", "REL", "REL"],
-    "weight": [0.0, 0.0, 0.1, 42.0]
+    {
+        "sourceNodeId": [0, 1, 2, 3],
+        "targetNodeId": [1, 2, 3, 0],
+        "relationshipType": ["REL", "REL", "REL", "REL"],
+        "weight": [0.0, 0.0, 0.1, 42.0]
+    }
 )
 
 G = gds.alpha.graph.construct(

graphdatascience/query_runner/cypher_graph_constructor.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import List, Set
+from typing import Any, List, Set, Tuple
 
 from pandas.core.frame import DataFrame
 
@@ -40,15 +40,18 @@ def run(self, node_dfs: List[DataFrame], relationship_dfs: List[DataFrame]) -> N
             "{readConcurrency: $read_concurrency, parameters: { nodes: $nodes, relationships: $relationships }})"
         )
 
+        node_query, nodes = self._node_query(node_dfs[0])
+        relationship_query, relationships = self._relationship_query(relationship_dfs[0])
+
         self._query_runner.run_query(
             query,
             {
                 "graph_name": self._graph_name,
-                "node_query": self._node_query(node_dfs[0]),
-                "relationship_query": self._relationship_query(relationship_dfs[0]),
+                "node_query": node_query,
+                "relationship_query": relationship_query,
                 "read_concurrency": self._concurrency,
-                "nodes": node_dfs[0].values.tolist(),
-                "relationships": relationship_dfs[0].values.tolist(),
+                "nodes": nodes,
+                "relationships": relationships,
             },
         )
 
@@ -59,23 +62,32 @@ def _is_enterprise(self) -> bool:
 
         return license == "Licensed"
 
-    def _node_query(self, node_df: DataFrame) -> str:
+    def _node_query(self, node_df: DataFrame) -> Tuple[str, List[List[Any]]]:
+        node_list = node_df.values.tolist()
         node_id_index = node_df.columns.get_loc("nodeId")
 
         label_query = ""
         if "labels" in node_df.keys():
             label_index = node_df.columns.get_loc("labels")
             label_query = f", node[{label_index}] as labels"
 
+            # Make sure every node has a list of labels
+            for node in node_list:
+                labels = node[label_index]
+                if isinstance(labels, List):
+                    continue
+                node[label_index] = [labels]
+
         property_query = ""
         property_columns: Set[str] = set(node_df.keys()) - {"nodeId", "labels"}
         if len(property_columns) > 0:
             property_queries = (f", node[{node_df.columns.get_loc(col)}] as {col}" for col in property_columns)
             property_query = "".join(property_queries)
 
-        return f"UNWIND $nodes as node RETURN node[{node_id_index}] as id{label_query}{property_query}"
+        return f"UNWIND $nodes as node RETURN node[{node_id_index}] as id{label_query}{property_query}", node_list
 
-    def _relationship_query(self, rel_df: DataFrame) -> str:
+    def _relationship_query(self, rel_df: DataFrame) -> Tuple[str, List[List[Any]]]:
+        rel_list = rel_df.values.tolist()
         source_id_index = rel_df.columns.get_loc("sourceNodeId")
         target_id_index = rel_df.columns.get_loc("targetNodeId")
 
@@ -93,5 +105,6 @@ def _relationship_query(self, rel_df: DataFrame) -> str:
         return (
             "UNWIND $relationships as relationship "
             f"RETURN relationship[{source_id_index}] as source, relationship[{target_id_index}] as target"
-            f"{type_query}{property_query}"
+            f"{type_query}{property_query}",
+            rel_list,
         )

graphdatascience/tests/integration/test_graph_ops.py

@@ -357,7 +357,7 @@ def test_graph_alpha_construct_without_arrow(gds_without_arrow: GraphDataScience
     nodes = pandas.DataFrame(
         {
             "nodeId": [0, 1, 2, 3],
-            "labels": [["A"], ["B"], ["C"], ["D"]],
+            "labels": [["A"], "B", ["C", "A"], ["D"]],
             "propA": [1337, 42, 8, 133742],
             "propB": [1338, 43, 9, 133743],
         }
@@ -377,6 +377,8 @@ def test_graph_alpha_construct_without_arrow(gds_without_arrow: GraphDataScience
     assert G.name() == "hello"
     assert G.node_count() == 4
     assert G.relationship_count() == 4
+    assert set(G.node_labels()) == {"A", "B", "C", "D"}
+    assert set(G.relationship_types()) == {"REL", "REL2"}
     assert set(G.node_properties("A")) == {"propA", "propB"}
     assert set(G.relationship_properties("REL")) == {"relPropA", "relPropB"}