Add expand option to Result.to_df.

robsdedude · robsdedude · commit dd015a4cda1e · 2022-03-15T19:39:10.000+01:00
This option (when `True`) will make the driver flatten nodes, relationships,
lists, and dicts into multiple columns of the DataFrame.
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -987,7 +987,7 @@ Path           :class:`neo4j.graph.Path`
 Node
 ====
 
-.. autoclass:: neo4j.graph.Node()
+.. autoclass:: neo4j.graph.Node
 
     .. describe:: node == other
 
@@ -1022,6 +1022,8 @@ Node
 
     .. autoattribute:: id
 
+    .. autoattribute:: element_id
+
     .. autoattribute:: labels
 
     .. automethod:: get
@@ -1036,7 +1038,7 @@ Node
 Relationship
 ============
 
-.. autoclass:: neo4j.graph.Relationship()
+.. autoclass:: neo4j.graph.Relationship
 
     .. describe:: relationship == other
 
@@ -1076,6 +1078,8 @@ Relationship
 
     .. autoattribute:: id
 
+    .. autoattribute:: element_id
+
     .. autoattribute:: nodes
 
     .. autoattribute:: start_node
@@ -1097,7 +1101,7 @@ Relationship
 Path
 ====
 
-.. autoclass:: neo4j.graph.Path()
+.. autoclass:: neo4j.graph.Path
 
     .. describe:: path == other
 
diff --git a/neo4j/_async/work/result.py b/neo4j/_async/work/result.py
@@ -20,7 +20,10 @@
 from warnings import warn
 
 from ..._async_compat.util import AsyncUtil
-from ...data import DataDehydrator
+from ...data import (
+    DataDehydrator,
+    RecordTableRowExporter,
+)
 from ...exceptions import (
     ResultConsumedError,
     ResultNotSingleError,
@@ -524,15 +527,74 @@ async def data(self, *keys):
 
     @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.
+    async def to_df(self, expand=False):
+        r"""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.
 
+        :param expand: if :const:`True`, some structures in the result will be
+            recursively expanded (flattened out into multiple columns) like so:
+
+            * :class:`.Node` objects under any variable ``<n>`` will be
+              expanded into columns (the recursion stops here)
+
+              * ``<n>().prop.<property_name>`` (any) for each property of the
+                node.
+              * ``<n>().element_id`` (str) the node's element id.
+                See :attr:`.Node.element_id`.
+              * ``<n>().labels`` (list of str) the node's labels.
+                See :attr:`.Node.labels`.
+
+            * :class:`.Relationship` objects under any variable ``<r>``
+              will be expanded into columns (the recursion stops here)
+
+              * ``<r>->.prop.<property_name>`` (any) for each property of the
+                relationship.
+              * ``<r>->.element_id`` (str) the relationship's element id.
+                See :attr:`.Relationship.element_id`.
+              * ``<r>->.start.element_id`` (str) the relationship's
+                start node's element id.
+                See :attr:`.Relationship.start_node`.
+              * ``<r>->.end.element_id`` (str) the relationship's
+                end node's element id.
+                See :attr:`.Relationship.end_node`.
+              * ``<r>->.type`` (str) the relationship's type.
+                See :attr:`.Relationship.type`.
+
+            * :const:`list` objects under any variable ``<l>`` will be expanded
+              into
+
+              * ``<l>[].0`` (any) the 1st list element
+              * ``<l>[].1`` (any) the 2nd list element
+              * ...
+
+            * :const:`dict` objects under any variable ``<d>`` will be expanded
+              into
+
+              * ``<d>{}.<key1>`` (any) the 1st key of the dict
+              * ``<d>{}.<key2>`` (any) the 2nd key of the dict
+              * ...
+
+            * :const:`list` and :const:`dict` objects are expanded recursively.
+              Example::
+
+                [{"foo": "bar", "baz": [42, 0]}, "foobar"]
+
+              will be expanded to::
+
+                {"0.foo": "bar", "0.baz.0": 42, "0.baz.1": 0, "1": "foobar"}
+
+            * Everything else (including :class:`.Path` objects) will not
+              be flattened.
+
+            :const:`dict` keys and variable names that contain ``.``  or ``\``
+            will be escaped with a backslash (``\.`` and ``\\`` respectively).
+        :type expand: bool
+
         :rtype: :py:class:`pandas.DataFrame`
         :raises ImportError: if `pandas` library is not available.
         :raises ResultConsumedError: if the transaction from which this result
@@ -545,7 +607,34 @@ async def to_df(self):
         """
         import pandas as pd
 
-        return pd.DataFrame(await self.values(), columns=self._keys)
+        if not expand:
+            return pd.DataFrame(await self.values(), columns=self._keys)
+        else:
+            df_keys = None
+            rows = []
+            async for record in self:
+                row = RecordTableRowExporter().transform(dict(record.items()))
+                if df_keys == row.keys():
+                    rows.append(row.values())
+                elif df_keys is None:
+                    df_keys = row.keys()
+                    rows.append(row.values())
+                elif df_keys is False:
+                    rows.append(row)
+                else:
+                    # The rows have different keys. We need to pass a list
+                    # of dicts to pandas
+                    rows = [{k: v for k, v in zip(df_keys, r)} for r in rows]
+                    df_keys = False
+                    rows.append(row)
+            if df_keys is False:
+                return pd.DataFrame(rows)
+            else:
+                columns = df_keys or [
+                    k.replace(".", "\\.").replace("\\", "\\\\")
+                    for k in self._keys
+                ]
+                return pd.DataFrame(rows, columns=columns)
 
     def closed(self):
         """Return True if the result has been closed.
diff --git a/neo4j/_sync/work/result.py b/neo4j/_sync/work/result.py
@@ -20,7 +20,10 @@
 from warnings import warn
 
 from ..._async_compat.util import Util
-from ...data import DataDehydrator
+from ...data import (
+    DataDehydrator,
+    RecordTableRowExporter,
+)
 from ...exceptions import (
     ResultConsumedError,
     ResultNotSingleError,
@@ -524,15 +527,74 @@ def data(self, *keys):
 
     @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.
+    def to_df(self, expand=False):
+        r"""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.
 
+        :param expand: if :const:`True`, some structures in the result will be
+            recursively expanded (flattened out into multiple columns) like so:
+
+            * :class:`.Node` objects under any variable ``<n>`` will be
+              expanded into columns (the recursion stops here)
+
+              * ``<n>().prop.<property_name>`` (any) for each property of the
+                node.
+              * ``<n>().element_id`` (str) the node's element id.
+                See :attr:`.Node.element_id`.
+              * ``<n>().labels`` (list of str) the node's labels.
+                See :attr:`.Node.labels`.
+
+            * :class:`.Relationship` objects under any variable ``<r>``
+              will be expanded into columns (the recursion stops here)
+
+              * ``<r>->.prop.<property_name>`` (any) for each property of the
+                relationship.
+              * ``<r>->.element_id`` (str) the relationship's element id.
+                See :attr:`.Relationship.element_id`.
+              * ``<r>->.start.element_id`` (str) the relationship's
+                start node's element id.
+                See :attr:`.Relationship.start_node`.
+              * ``<r>->.end.element_id`` (str) the relationship's
+                end node's element id.
+                See :attr:`.Relationship.end_node`.
+              * ``<r>->.type`` (str) the relationship's type.
+                See :attr:`.Relationship.type`.
+
+            * :const:`list` objects under any variable ``<l>`` will be expanded
+              into
+
+              * ``<l>[].0`` (any) the 1st list element
+              * ``<l>[].1`` (any) the 2nd list element
+              * ...
+
+            * :const:`dict` objects under any variable ``<d>`` will be expanded
+              into
+
+              * ``<d>{}.<key1>`` (any) the 1st key of the dict
+              * ``<d>{}.<key2>`` (any) the 2nd key of the dict
+              * ...
+
+            * :const:`list` and :const:`dict` objects are expanded recursively.
+              Example::
+
+                [{"foo": "bar", "baz": [42, 0]}, "foobar"]
+
+              will be expanded to::
+
+                {"0.foo": "bar", "0.baz.0": 42, "0.baz.1": 0, "1": "foobar"}
+
+            * Everything else (including :class:`.Path` objects) will not
+              be flattened.
+
+            :const:`dict` keys and variable names that contain ``.``  or ``\``
+            will be escaped with a backslash (``\.`` and ``\\`` respectively).
+        :type expand: bool
+
         :rtype: :py:class:`pandas.DataFrame`
         :raises ImportError: if `pandas` library is not available.
         :raises ResultConsumedError: if the transaction from which this result
@@ -545,7 +607,34 @@ def to_df(self):
         """
         import pandas as pd
 
-        return pd.DataFrame(self.values(), columns=self._keys)
+        if not expand:
+            return pd.DataFrame(self.values(), columns=self._keys)
+        else:
+            df_keys = None
+            rows = []
+            for record in self:
+                row = RecordTableRowExporter().transform(dict(record.items()))
+                if df_keys == row.keys():
+                    rows.append(row.values())
+                elif df_keys is None:
+                    df_keys = row.keys()
+                    rows.append(row.values())
+                elif df_keys is False:
+                    rows.append(row)
+                else:
+                    # The rows have different keys. We need to pass a list
+                    # of dicts to pandas
+                    rows = [{k: v for k, v in zip(df_keys, r)} for r in rows]
+                    df_keys = False
+                    rows.append(row)
+            if df_keys is False:
+                return pd.DataFrame(rows)
+            else:
+                columns = df_keys or [
+                    k.replace(".", "\\.").replace("\\", "\\\\")
+                    for k in self._keys
+                ]
+                return pd.DataFrame(rows, columns=columns)
 
     def closed(self):
         """Return True if the result has been closed.
diff --git a/neo4j/data.py b/neo4j/data.py
@@ -297,6 +297,59 @@ def transform(self, x):
             return x
 
 
+class RecordTableRowExporter(DataTransformer):
+    """Transformer class used by the :meth:`.Result.to_df` method."""
+
+    def transform(self, x):
+        assert isinstance(x, Mapping)
+        t = type(x)
+        return t(item
+                 for k, v in x.items()
+                 for item in self._transform(
+                     v, prefix=k.replace("\\", "\\\\").replace(".", "\\.")
+                 ).items())
+
+    def _transform(self, x, prefix):
+        if isinstance(x, Node):
+            res = {
+                "%s().element_id" % prefix: x.element_id,
+                "%s().labels" % prefix: x.labels,
+            }
+            res.update(("%s().prop.%s" % (prefix, k), v) for k, v in x.items())
+            return res
+        elif isinstance(x, Relationship):
+            res = {
+                "%s->.element_id" % prefix: x.element_id,
+                "%s->.start.element_id" % prefix: x.start_node.element_id,
+                "%s->.end.element_id" % prefix: x.end_node.element_id,
+                "%s->.type" % prefix: x.__class__.__name__,
+            }
+            res.update(("%s->.prop.%s" % (prefix, k), v) for k, v in x.items())
+            return res
+        elif isinstance(x, Path) or isinstance(x, str):
+            return {prefix: x}
+        elif isinstance(x, Sequence):
+            return dict(
+                item
+                for i, v in enumerate(x)
+                for item in self._transform(
+                    v, prefix="%s[].%i" % (prefix, i)
+                ).items()
+            )
+        elif isinstance(x, Mapping):
+            t = type(x)
+            return t(
+                item
+                for k, v in x.items()
+                for item in self._transform(
+                    v, prefix="%s{}.%s" % (prefix, k.replace("\\", "\\\\")
+                                                    .replace(".", "\\."))
+                ).items()
+            )
+        else:
+            return {prefix: x}
+
+
 class DataHydrator:
     # TODO: extend DataTransformer
 
diff --git a/neo4j/graph/__init__.py b/neo4j/graph/__init__.py
@@ -207,6 +207,10 @@ def id(self):
         Depending on the version of the server this entity was retrieved from,
         this may be empty (None).
 
+        .. Warning::
+            This value can change for the same entity across multiple
+            queries. Don't rely on it for cross-query computations.
+
         .. deprecated:: 5.0
             Use :attr:`.element_id` instead.
 
@@ -218,7 +222,11 @@ def id(self):
     def element_id(self):
         """The identity of this entity in its container :class:`.Graph`.
 
-        .. added:: 5.0
+        .. Warning::
+            This value can change for the same entity across multiple
+            queries. Don't rely on it for cross-query computations.
+
+        .. versionadded:: 5.0
 
         :rtype: str
         """
diff --git a/tests/unit/async_/work/test_result.py b/tests/unit/async_/work/test_result.py
diff --git a/tests/unit/sync/work/test_result.py b/tests/unit/sync/work/test_result.py
