Add Result.to_df to export records as pandas DataFrame

Author: robsdedude

docs/source/api.rst

@@ -808,14 +808,14 @@ A :class:`neo4j.Result` is attached to an active connection, through a :class:`n
 
     .. automethod:: graph
 
-       **This is experimental.** (See :ref:`filter-warnings-ref`)
-
     .. automethod:: value
 
     .. automethod:: values
 
     .. automethod:: data
 
+    .. automethod:: to_df
+
     .. automethod:: closed
 
 See https://neo4j.com/docs/python-manual/current/cypher-workflow/#python-driver-type-mapping for more about type mapping.

docs/source/async_api.rst

@@ -502,14 +502,14 @@ A :class:`neo4j.AsyncResult` is attached to an active connection, through a :cla
 
     .. automethod:: graph
 
-       **This is experimental.** (See :ref:`filter-warnings-ref`)
-
     .. automethod:: value
 
     .. automethod:: values
 
     .. automethod:: data
 
+    .. automethod:: to_df
+
     .. automethod:: closed
 
 See https://neo4j.com/docs/python-manual/current/cypher-workflow/#python-driver-type-mapping for more about type mapping.

neo4j/_async/work/result.py

@@ -24,6 +24,7 @@
     ResultConsumedError,
     ResultNotSingleError,
 )
+from ...meta import experimental
 from ...work import ResultSummary
 from ..io import ConnectionErrorHandler
 
@@ -392,6 +393,8 @@ async def graph(self):
         :raises ResultConsumedError: if the transaction from which this result
             was obtained has been closed or the Result has been explicitly
             consumed.
+
+        **This is experimental.** (See :ref:`filter-warnings-ref`)
         """
         await self._buffer_all()
         return self._hydrant.graph
@@ -444,6 +447,28 @@ async def data(self, *keys):
         """
         return [record.data(*keys) async for record in self]
 
+    @experimental("pandas support is experimental and might be changed or "
+                  "removed in future versions")
+    async def to_df(self):
+        """Convert (the rest of) the result to a pandas DataFrame.
+
+        This method is only available if the `pandas` library is installed.
+
+        ``tx.run("UNWIND range(1, 10) AS n RETURN n, n+1 as m").to_df()``, for
+        instance will return a DataFrame with two columns: ``n`` and ``m`` and
+        10 rows.
+
+        :rtype: :py:class:`pandas.DataFrame`
+        :raises ImportError: if `pandas` library is not available.
+
+        **This is experimental.**
+        ``pandas`` support might be changed or removed in future versions
+        without warning. (See :ref:`filter-warnings-ref`)
+        """
+        import pandas as pd
+
+        return pd.DataFrame(await self.values(), columns=self._keys)
+
     def closed(self):
         """Return True if the result has been closed.
 

neo4j/_sync/work/result.py

@@ -24,6 +24,7 @@
     ResultConsumedError,
     ResultNotSingleError,
 )
+from ...meta import experimental
 from ...work import ResultSummary
 from ..io import ConnectionErrorHandler
 
@@ -392,6 +393,8 @@ def graph(self):
         :raises ResultConsumedError: if the transaction from which this result
             was obtained has been closed or the Result has been explicitly
             consumed.
+
+        **This is experimental.** (See :ref:`filter-warnings-ref`)
         """
         self._buffer_all()
         return self._hydrant.graph
@@ -444,6 +447,28 @@ def data(self, *keys):
         """
         return [record.data(*keys) for record in self]
 
+    @experimental("pandas support is experimental and might be changed or "
+                  "removed in future versions")
+    def to_df(self):
+        """Convert (the rest of) the result to a pandas DataFrame.
+
+        This method is only available if the `pandas` library is installed.
+
+        ``tx.run("UNWIND range(1, 10) AS n RETURN n, n+1 as m").to_df()``, for
+        instance will return a DataFrame with two columns: ``n`` and ``m`` and
+        10 rows.
+
+        :rtype: :py:class:`pandas.DataFrame`
+        :raises ImportError: if `pandas` library is not available.
+
+        **This is experimental.**
+        ``pandas`` support might be changed or removed in future versions
+        without warning. (See :ref:`filter-warnings-ref`)
+        """
+        import pandas as pd
+
+        return pd.DataFrame(self.values(), columns=self._keys)
+
     def closed(self):
         """Return True if the result has been closed.
 

tests/requirements.txt

@@ -7,3 +7,4 @@ pytest-cov>=3.0.0
 pytest-mock>=3.6.1
 mock>=4.0.3
 teamcity-messages>=1.29
+pandas>=1.4.1

tests/unit/async_/work/test_result.py

@@ -18,6 +18,7 @@
 
 from unittest import mock
 
+import pandas as pd
 import pytest
 
 from neo4j import (
@@ -30,18 +31,21 @@
     Version,
 )
 from neo4j._async_compat.util import AsyncUtil
-from neo4j.data import DataHydrator
+from neo4j.data import (
+    DataHydrator,
+    Node,
+)
 from neo4j.exceptions import ResultNotSingleError
+from neo4j.packstream import Structure
 
 from ...._async_compat import mark_async_test
 
 
 class Records:
     def __init__(self, fields, records):
-        assert all(len(fields) == len(r) for r in records)
-        self.fields = fields
-        # self.records = [{"record_values": r} for r in records]
-        self.records = records
+        self.fields = tuple(fields)
+        self.records = tuple(records)
+        assert all(len(self.fields) == len(r) for r in self.records)
 
     def __len__(self):
         return self.records.__len__()
@@ -469,3 +473,46 @@ async def test_data(num_records):
     assert await result.data("hello", "world") == expected_data
     for record in records:
         assert record.data.called_once_with("hello", "world")
+
+
+@pytest.mark.parametrize(
+    ("keys", "values", "types", "instances"),
+    (
+        (["i"], zip(range(5)), ["int64"], None),
+        (["x"], zip((n - .5) / 5 for n in range(5)), ["float64"], None),
+        (["s"], zip(("foo", "bar", "baz", "foobar")), ["object"], None),
+        (["l"], zip(([1, 2], [3, 4])), ["object"], None),
+        (
+            ["n"],
+            zip((
+                Structure(b"N", 0, ["LABEL_A"], {"a": 1, "b": 2}),
+                Structure(b"N", 2, ["LABEL_B"], {"a": 1, "c": 1.2}),
+                Structure(b"N", 1, ["LABEL_A", "LABEL_B"], {"a": [1, "a"]}),
+            )),
+            ["object"],
+            [Node]
+        ),
+    )
+)
+@mark_async_test
+async def test_to_df(keys, values, types, instances):
+    values = list(values)
+    connection = AsyncConnectionStub(records=Records(keys, values))
+    result = AsyncResult(connection, DataHydrator(), 1, noop, noop)
+    await result._run("CYPHER", {}, None, None, "r", None)
+    df = await result.to_df()
+
+    assert isinstance(df, pd.DataFrame)
+    assert df.keys().to_list() == keys
+    assert len(df) == len(values)
+    assert df.dtypes.to_list() == types
+
+    expected_df = pd.DataFrame(
+        {k: [v[i] for v in values] for i, k in enumerate(keys)}
+    )
+
+    if instances:
+        for i, k in enumerate(keys):
+            assert all(isinstance(v, instances[i]) for v in df[k])
+    else:
+        assert df.equals(expected_df)

tests/unit/sync/work/test_result.py

@@ -18,6 +18,7 @@
 
 from unittest import mock
 
+import pandas as pd
 import pytest
 
 from neo4j import (
@@ -30,18 +31,21 @@
     Version,
 )
 from neo4j._async_compat.util import Util
-from neo4j.data import DataHydrator
+from neo4j.data import (
+    DataHydrator,
+    Node,
+)
 from neo4j.exceptions import ResultNotSingleError
+from neo4j.packstream import Structure
 
 from ...._async_compat import mark_sync_test
 
 
 class Records:
     def __init__(self, fields, records):
-        assert all(len(fields) == len(r) for r in records)
-        self.fields = fields
-        # self.records = [{"record_values": r} for r in records]
-        self.records = records
+        self.fields = tuple(fields)
+        self.records = tuple(records)
+        assert all(len(self.fields) == len(r) for r in self.records)
 
     def __len__(self):
         return self.records.__len__()
@@ -469,3 +473,46 @@ def test_data(num_records):
     assert result.data("hello", "world") == expected_data
     for record in records:
         assert record.data.called_once_with("hello", "world")
+
+
+@pytest.mark.parametrize(
+    ("keys", "values", "types", "instances"),
+    (
+        (["i"], zip(range(5)), ["int64"], None),
+        (["x"], zip((n - .5) / 5 for n in range(5)), ["float64"], None),
+        (["s"], zip(("foo", "bar", "baz", "foobar")), ["object"], None),
+        (["l"], zip(([1, 2], [3, 4])), ["object"], None),
+        (
+            ["n"],
+            zip((
+                Structure(b"N", 0, ["LABEL_A"], {"a": 1, "b": 2}),
+                Structure(b"N", 2, ["LABEL_B"], {"a": 1, "c": 1.2}),
+                Structure(b"N", 1, ["LABEL_A", "LABEL_B"], {"a": [1, "a"]}),
+            )),
+            ["object"],
+            [Node]
+        ),
+    )
+)
+@mark_sync_test
+def test_to_df(keys, values, types, instances):
+    values = list(values)
+    connection = ConnectionStub(records=Records(keys, values))
+    result = Result(connection, DataHydrator(), 1, noop, noop)
+    result._run("CYPHER", {}, None, None, "r", None)
+    df = result.to_df()
+
+    assert isinstance(df, pd.DataFrame)
+    assert df.keys().to_list() == keys
+    assert len(df) == len(values)
+    assert df.dtypes.to_list() == types
+
+    expected_df = pd.DataFrame(
+        {k: [v[i] for v in values] for i, k in enumerate(keys)}
+    )
+
+    if instances:
+        for i, k in enumerate(keys):
+            assert all(isinstance(v, instances[i]) for v in df[k])
+    else:
+        assert df.equals(expected_df)