fix the to_df in AlgorithmResult and Edges (#1820)

* fix the `to_df` in AlgorithmResult such that it returns nodes instead of useless internal ids

* remove PyDirection and make Direction directly constructable from string in python

* fix the dataframe test

* return node id instead of view

* detect exploded edges in to_df

* add test to check that calling `edges.to_df(explode=True)` is the same as `edges.explode().to_df()`

* fix docstring

Author: ljeub-pometry

python/python/raphtory/__init__.pyi

@@ -100,12 +100,12 @@ class AlgorithmResult:
             A sorted vector of tuples containing keys of type `H` and values of type `Y`.
         """
 
-    def to_df(self):
+    def to_df(self) -> DataFrame:
         """
         Creates a dataframe from the result
 
         Returns:
-            A `pandas.DataFrame` containing the result
+            DataFrame: A `pandas.DataFrame` containing the result
         """
 
     def to_string(self):
@@ -5656,14 +5656,6 @@ class PropertyFilter:
     def __init__(self):
         """Initialize self.  See help(type(self)) for accurate signature."""
 
-class PyDirection:
-    """A direction used by an edge, being incoming or outgoing"""
-
-    def __init__(self, direction):
-        """Initialize self.  See help(type(self)) for accurate signature."""
-
-    def as_str(self): ...
-
 class PyGraphEncoder:
     def __init__(self):
         """Initialize self.  See help(type(self)) for accurate signature."""

python/python/raphtory/algorithms/__init__.pyi

@@ -46,7 +46,7 @@ def average_degree(g: GraphView):
 def balance(
     g: GraphView,
     name: str = "weight",
-    direction: Direction = "BOTH",
+    direction: Direction = "both",
     threads: Optional[int] = None,
 ) -> AlgorithmResult:
     """
@@ -57,7 +57,7 @@ def balance(
     Arguments:
         g (GraphView): The graph view on which the operation is to be performed.
         name (str): The name of the edge property used as the weight. Defaults to "weight".
-        direction (Direction): Specifies the direction of the edges to be considered for summation. Defaults to "BOTH".
+        direction (Direction): Specifies the direction of the edges to be considered for summation. Defaults to "both".
                 * "OUT": Only consider outgoing edges.
                 * "IN": Only consider incoming edges.
                 * "BOTH": Consider both outgoing and incoming edges. This is the default.
@@ -108,7 +108,7 @@ def dijkstra_single_source_shortest_paths(
     g: GraphView,
     source: InputNode,
     targets: list[InputNode],
-    direction: Direction = "BOTH",
+    direction: Direction = "both",
     weight: str = "weight",
 ) -> dict:
     """
@@ -118,7 +118,7 @@ def dijkstra_single_source_shortest_paths(
         g (GraphView): The graph to search in.
         source (InputNode): The source node.
         targets (list[InputNode]): A list of target nodes.
-        direction (Direction): The direction of the edges to be considered for the shortest path. Defaults to "BOTH".
+        direction (Direction): The direction of the edges to be considered for the shortest path. Defaults to "both".
         weight (str): The name of the weight property for the edges. Defaults to "weight".
 
     Returns:

python/python/raphtory/typing.py

@@ -13,7 +13,7 @@
     dict[str, "Prop"],
 ]
 
-Direction = Literal["IN", "OUT", "BOTH"]
+Direction = Literal["in", "out", "both"]
 
 InputNode = Union[int, str, "Node"]
 

python/tests/test_algorithms.py

@@ -1,7 +1,7 @@
 import pytest
 import pandas as pd
 import pandas.core.frame
-from raphtory import Graph, PersistentGraph, PyDirection
+from raphtory import Graph, PersistentGraph
 from raphtory import algorithms
 from raphtory import graph_loader
 
@@ -223,8 +223,8 @@ def test_algo_result():
     assert len(actual.group_by()[1]) == 8
     assert type(actual.to_df()) == pandas.core.frame.DataFrame
     df = actual.to_df()
-    expected_result = pd.DataFrame({"Key": [1], "Value": [1]})
-    row_with_one = df[df["Key"] == 1]
+    expected_result = pd.DataFrame({"Node": 1, "Value": [1]})
+    row_with_one = df[df["Node"] == 1]
     row_with_one.reset_index(inplace=True, drop=True)
     assert row_with_one.equals(expected_result)
     # Algo Str u64
@@ -254,7 +254,7 @@ def test_algo_result():
         "8": 0.11786468661230831,
     }
     assert actual.get_all_with_names() == expected_result
-    assert actual.get("Not a node") == None
+    assert actual.get("Not a node") is None
     assert len(actual.to_df()) == 8
     # algo str vector
     actual = algorithms.temporally_reachable_nodes(g, 20, 11, [1, 2], [4, 5])
@@ -463,19 +463,13 @@ def test_balance_algorithm():
     ]
     for src, dst, val, time in edges_str:
         g.add_edge(time, src, dst, {"value_dec": val})
-    result = algorithms.balance(
-        g, "value_dec", PyDirection("BOTH"), None
-    ).get_all_with_names()
+    result = algorithms.balance(g, "value_dec", "both", None).get_all_with_names()
     assert result == {"1": -26.0, "2": 7.0, "3": 12.0, "4": 5.0, "5": 2.0}
 
-    result = algorithms.balance(
-        g, "value_dec", PyDirection("IN"), None
-    ).get_all_with_names()
+    result = algorithms.balance(g, "value_dec", "in", None).get_all_with_names()
     assert result == {"1": 6.0, "2": 12.0, "3": 15.0, "4": 20.0, "5": 2.0}
 
-    result = algorithms.balance(
-        g, "value_dec", PyDirection("OUT"), None
-    ).get_all_with_names()
+    result = algorithms.balance(g, "value_dec", "out", None).get_all_with_names()
     assert result == {"1": -32.0, "2": -5.0, "3": -3.0, "4": -15.0, "5": 0.0}
 
 

python/tests/test_graph_conversions.py

@@ -1124,6 +1124,11 @@ def test_to_df():
         pd.read_json(base_dir / "expected/dataframe_output/edge_df_explode.json"),
     )
 
+    compare_df(
+        g.edges.explode().to_df(),
+        pd.read_json(base_dir / "expected/dataframe_output/edge_df_explode.json"),
+    )
+
     # compare_df(
     #     g.edges.to_df(explode=True, convert_datetime=True),
     #     pd.read_json(base_dir / "expected/dataframe_output/edge_df_explode_datetime.json")

python/tests/test_graphdb/test_graphdb.py

@@ -7,7 +7,7 @@
 import pandas as pd
 import pandas.core.frame
 import pytest
-from raphtory import Graph, PersistentGraph, PyDirection
+from raphtory import Graph, PersistentGraph
 from raphtory import algorithms
 from raphtory import graph_loader
 import tempfile

python/tests/test_graphdb/test_iterables.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import pandas.core.frame
 import pytest
-from raphtory import Graph, PersistentGraph, PyDirection
+from raphtory import Graph, PersistentGraph
 from raphtory import algorithms
 from raphtory import graph_loader
 import tempfile

python/tests/test_loaders/test_load_from_pandas.py

@@ -6,7 +6,7 @@
 import pandas as pd
 import pandas.core.frame
 import pytest
-from raphtory import Graph, PersistentGraph, PyDirection
+from raphtory import Graph, PersistentGraph
 from raphtory import algorithms
 from raphtory import graph_loader
 import tempfile

raphtory-api/src/python/direction.rs

@@ -0,0 +1,16 @@
+use crate::core::Direction;
+use pyo3::{exceptions::PyTypeError, FromPyObject, PyAny, PyResult};
+
+impl<'source> FromPyObject<'source> for Direction {
+    fn extract(ob: &'source PyAny) -> PyResult<Self> {
+        let value: &str = ob.extract()?;
+        match value {
+            "out" => Ok(Direction::OUT),
+            "in" => Ok(Direction::IN),
+            "both" => Ok(Direction::BOTH),
+            _ => Err(PyTypeError::new_err(PyTypeError::new_err(
+                "Direction must be one of { 'out', 'in', 'both' }",
+            ))),
+        }
+    }
+}

raphtory-api/src/python/mod.rs

@@ -1,2 +1,3 @@
 mod arcstr;
+mod direction;
 mod gid;