Make sure to return the correct columns for stream[Node/Relationship]Properties

Author: DarthMax

graphdatascience/graph/graph_proc_runner.py

@@ -113,7 +113,7 @@ def streamNodeProperties(
 
         # new format was requested, but the query was run via Cypher
         if separate_property_columns and "propertyValue" in result.keys():
-            return result.pivot_table("propertyValue", "nodeId", columns="nodeProperty")
+            return result.pivot_table("propertyValue", "nodeId", columns="nodeProperty").reset_index()
         # old format was requested but the query was run via Arrow
         elif not separate_property_columns and "propertyValue" not in result.keys():
             return result.melt(id_vars=["nodeId"]).rename(
@@ -149,7 +149,7 @@ def streamRelationshipProperties(
         if separate_property_columns and "propertyValue" in result.keys():
             return result.pivot_table(
                 "propertyValue", ["sourceNodeId", "targetNodeId", "relationshipType"], columns="relationshipProperty"
-            )
+            ).reset_index()
         # old format was requested but the query was run via Arrow
         elif not separate_property_columns and "propertyValue" not in result.keys():
             return result.melt(id_vars=["sourceNodeId", "targetNodeId", "relationshipType"]).rename(

graphdatascience/tests/integration/test_graph_ops.py

@@ -168,17 +168,20 @@ def test_graph_streamNodeProperties(gds: GraphDataScience) -> None:
 
     result = gds.graph.streamNodeProperties(G, ["x", "y"], concurrency=2)
 
+    assert list(result.keys()) == ["nodeId", "nodeProperty", "propertyValue"]
+
     x_values = result[result.nodeProperty == "x"]
     assert {e for e in x_values["propertyValue"]} == {1, 2, 3}
 
     y_values = result[result.nodeProperty == "y"]
     assert {e for e in y_values["propertyValue"]} == {2, 3, 4}
 
 
-def test_graph_streamNodeProperties_merge_property_columns(gds: GraphDataScience) -> None:
+def test_graph_streamNodeProperties_separate_property_columns(gds: GraphDataScience) -> None:
     G, _ = gds.graph.project(GRAPH_NAME, {"Node": {"properties": ["x", "y"]}}, "*")
 
     result = gds.graph.streamNodeProperties(G, ["x", "y"], separate_property_columns=True, concurrency=2)
+    assert list(result.keys()) == ["nodeId", "x", "y"]
     assert {e for e in result["x"]} == {1, 2, 3}
     assert {e for e in result["y"]} == {2, 3, 4}
 
@@ -187,17 +190,24 @@ def test_graph_streamNodeProperties_without_arrow(gds_without_arrow: GraphDataSc
     G, _ = gds_without_arrow.graph.project(GRAPH_NAME, {"Node": {"properties": ["x", "y"]}}, "*")
 
     result = gds_without_arrow.graph.streamNodeProperties(G, ["x", "y"], concurrency=2)
+
+    assert list(result.keys()) == ["nodeId", "nodeProperty", "propertyValue"]
+
     x_values = result[result.nodeProperty == "x"]
     assert {e for e in x_values["propertyValue"]} == {1, 2, 3}
 
     y_values = result[result.nodeProperty == "y"]
     assert {e for e in y_values["propertyValue"]} == {2, 3, 4}
 
 
-def test_graph_streamNodeProperties_without_arrow_merge_property_columns(gds_without_arrow: GraphDataScience) -> None:
+def test_graph_streamNodeProperties_without_arrow_separate_property_columns(
+    gds_without_arrow: GraphDataScience,
+) -> None:
     G, _ = gds_without_arrow.graph.project(GRAPH_NAME, {"Node": {"properties": ["x", "y"]}}, "*")
 
     result = gds_without_arrow.graph.streamNodeProperties(G, ["x", "y"], separate_property_columns=True, concurrency=2)
+
+    assert list(result.keys()) == ["nodeId", "x", "y"]
     assert {e for e in result["x"]} == {1, 2, 3}
     assert {e for e in result["y"]} == {2, 3, 4}
 
@@ -221,17 +231,26 @@ def test_graph_streamRelationshipProperties(gds: GraphDataScience) -> None:
 
     result = gds.graph.streamRelationshipProperties(G, ["relX", "relY"], concurrency=2)
 
+    assert list(result.keys()) == [
+        "sourceNodeId",
+        "targetNodeId",
+        "relationshipType",
+        "relationshipProperty",
+        "propertyValue",
+    ]
+
     x_values = result[result.relationshipProperty == "relX"]
     assert {e for e in x_values["propertyValue"]} == {4, 5, 6}
     y_values = result[result.relationshipProperty == "relY"]
     assert {e for e in y_values["propertyValue"]} == {5, 6, 7}
 
 
-def test_graph_streamRelationshipProperties_merge_property_columns(gds: GraphDataScience) -> None:
+def test_graph_streamRelationshipProperties_separate_property_columns(gds: GraphDataScience) -> None:
     G, _ = gds.graph.project(GRAPH_NAME, "*", {"REL": {"properties": ["relX", "relY"]}})
 
     result = gds.graph.streamRelationshipProperties(G, ["relX", "relY"], separate_property_columns=True, concurrency=2)
-    assert {e for e in result["relationshipType"]} == {"REL", "REL", "REL"}
+
+    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}
 
@@ -241,20 +260,30 @@ def test_graph_streamRelationshipProperties_without_arrow(gds_without_arrow: Gra
 
     result = gds_without_arrow.graph.streamRelationshipProperties(G, ["relX", "relY"], concurrency=2)
 
+    assert list(result.keys()) == [
+        "sourceNodeId",
+        "targetNodeId",
+        "relationshipType",
+        "relationshipProperty",
+        "propertyValue",
+    ]
+
     x_values = result[result.relationshipProperty == "relX"]
     assert {e for e in x_values["propertyValue"]} == {4, 5, 6}
     y_values = result[result.relationshipProperty == "relY"]
     assert {e for e in y_values["propertyValue"]} == {5, 6, 7}
 
 
-def test_graph_streamRelationshipProperties_without_arrow_merge_property_columns(
+def test_graph_streamRelationshipProperties_without_arrow_separate_property_columns(
     gds_without_arrow: GraphDataScience,
 ) -> None:
     G, _ = gds_without_arrow.graph.project(GRAPH_NAME, "*", {"REL": {"properties": ["relX", "relY"]}})
 
     result = gds_without_arrow.graph.streamRelationshipProperties(
         G, ["relX", "relY"], 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}