diff --git a/neo4j_arrow/_client.py b/neo4j_arrow/_client.py index 8f35dc8..6cee9d9 100644 --- a/neo4j_arrow/_client.py +++ b/neo4j_arrow/_client.py @@ -100,9 +100,9 @@ def _send_action(self, action: str, body: Dict[str, Any]) -> Dict[str, Any]: raise e @classmethod - def _nop(*args, **kwargs): + def _nop(cls, data: Arrow) -> Arrow: """Used as a no-op mapping function.""" - pass + return data @classmethod def _node_mapper(cls, model: Graph, source_field: Optional[str] = None): diff --git a/neo4j_arrow/test_client.py b/neo4j_arrow/test_client.py index c8add40..61db83b 100644 --- a/neo4j_arrow/test_client.py +++ b/neo4j_arrow/test_client.py @@ -4,6 +4,15 @@ import pyarrow as pa +def test_nop_mapper(): + mapper = Neo4jArrowClient._nop + t = pa.table({"labels": ["Junk"], "nodeId": ["junk_id"]}) + result = mapper(t) + assert result is not None + assert "nodeId" in result.schema.names + assert "labels" in result.schema.names + + def test_node_mapper(): SRC_KEY = "gcs_source" g = Graph(name="junk", nodes=[Node(source="gs://.*/junk.*parquet",