Warn when limiting with from_neo4j

Co-Authored-By: Florentin Dörre <[email protected]>

Author: adamnsch

python-wrapper/src/neo4j_viz/neo4j.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import Optional, Union
 
 import neo4j.graph
@@ -59,6 +60,15 @@ def from_neo4j(
     elif isinstance(data, neo4j.graph.Graph):
         graph = data
     elif isinstance(data, Driver):
+        rel_count = data.execute_query(
+            "MATCH ()-[r]->() RETURN count(r) as count",
+            routing_=RoutingControl.READ,
+            result_transformer_=Result.single,
+        ).get("count")  # type: ignore[union-attr]
+        if rel_count > row_limit:
+            warnings.warn(
+                f"Database relationship count ({rel_count}) exceeds `row_limit` ({row_limit}), so limiting will be applied. Increase the `row_limit` if needed"
+            )
         graph = data.execute_query(
             f"MATCH (n)-[r]->(m) RETURN n,r,m LIMIT {row_limit}",
             routing_=RoutingControl.READ,

python-wrapper/tests/test_neo4j.py

@@ -1,3 +1,4 @@
+import re
 from typing import Generator
 
 import neo4j
@@ -249,3 +250,16 @@ def test_from_neo4j_graph_driver(neo4j_session: Session, neo4j_driver: Driver) -
         (node_ids[0], node_ids[1], "KNOWS"),
         (node_ids[1], node_ids[0], "RELATED"),
     ]
+
+
+@pytest.mark.requires_neo4j_and_gds
+def test_from_neo4j_graph_row_limit_warning(neo4j_session: Session, neo4j_driver: Driver) -> None:
+    neo4j_session.run("MATCH (a:_CI_A|_CI_B)-[r]->(b) RETURN a, b, r ORDER BY a").graph()
+
+    with pytest.warns(
+        UserWarning,
+        match=re.escape(
+            "Database relationship count (2) exceeds `row_limit` (1), so limiting will be applied. Increase the `row_limit` if needed"
+        ),
+    ):
+        from_neo4j(neo4j_driver, row_limit=1)