Fix for relationshipType as str

Author: orazve

examples/kge-predict-transe-pyg-train.ipynb

@@ -307,13 +307,16 @@
    "outputs": [],
    "source": [
     "def create_data_from_graph(relationship_type):\n",
-    "    rels_tmp = gds.graph.relationshipProperties.stream(ttv_G, [\"rel_id\"], [relationship_type])\n",
+    "    rels_tmp = gds.graph.relationshipProperties.stream(\n",
+    "        ttv_G, [\"rel_id\"], relationship_type, separate_property_columns=True\n",
+    "    )\n",
+    "    # print(rels_tmp)\n",
     "    topology = [\n",
     "        rels_tmp.sourceNodeId.map(lambda x: nodeId_to_id[x]),\n",
     "        rels_tmp.targetNodeId.map(lambda x: nodeId_to_id[x]),\n",
     "    ]\n",
     "    edge_index = torch.tensor(topology, dtype=torch.long)\n",
-    "    edge_type = torch.tensor(rels_tmp.propertyValue.astype(int), dtype=torch.long)\n",
+    "    edge_type = torch.tensor(rels_tmp.rel_id.astype(int), dtype=torch.long)\n",
     "    data = Data(edge_index=edge_index, edge_type=edge_type)\n",
     "    data.num_nodes = len(nodeId_to_id)\n",
     "    display(data)\n",

graphdatascience/graph/base_graph_proc_runner.py

@@ -517,8 +517,7 @@ def writeRelationship(
         ).squeeze()
 
     @multimethod
-    def removeNodeProperties(self) -> None:
-        ...
+    def removeNodeProperties(self) -> None: ...
 
     @removeNodeProperties.register
     @graph_type_check

graphdatascience/graph/graph_entity_ops_runner.py

@@ -209,6 +209,8 @@ def stream(
     ) -> DataFrame:
         self._namespace += ".stream"
 
+        relationship_types = [relationship_types] if isinstance(relationship_types, str) else relationship_types
+
         result = self._handle_properties(G, relationship_properties, relationship_types, config)
 
         # new format was requested, but the query was run via Cypher

graphdatascience/tests/integration/test_graph_ops.py

@@ -704,6 +704,26 @@ def test_graph_relationshipProperties_stream_with_arrow_rel_as_str(gds: GraphDat
     assert {e for e in y_values["propertyValue"]} == {5, 6, 7}
 
 
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
+def test_graph_relationshipProperties_stream_with_arrow_rel_as_str_sep(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, "*", {"REL": {"properties": ["relX", "relY"]}})
+
+    result = gds.graph.relationshipProperties.stream(
+        G, ["relX", "relY"], "REL", separate_property_columns=True, concurrency=2
+    )
+
+    assert list(result.keys()) == [
+        "sourceNodeId",
+        "targetNodeId",
+        "relationshipType",
+        "relX",
+        "relY",
+    ]
+
+    assert {e for e in result["relX"]} == {4, 5, 6}
+    assert {e for e in result["relY"]} == {5, 6, 7}
+
+
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
 def test_graph_relationshipProperties_stream_with_arrow_rel_as_str_separate(gds: GraphDataScience) -> None:
     G, _ = gds.graph.project(GRAPH_NAME, "*", {"REL": {"properties": ["relX", "relY"]}})