Skip to content

Add parse_dates option to Result.to_df export method #716

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 25, 2022
Merged

Add parse_dates option to Result.to_df export method #716

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 30 additions & 5 deletions neo4j/_async/work/result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@
ResultNotSingleError,
)
from ...meta import experimental
from ...time import (
Date,
DateTime,
)
from ...work import ResultSummary
from ..io import ConnectionErrorHandler

Expand Down Expand Up @@ -527,7 +531,7 @@ async def data(self, *keys):

@experimental("pandas support is experimental and might be changed or "
"removed in future versions")
async def to_df(self, expand=False):
async def to_df(self, expand=False, parse_dates=False):
r"""Convert (the rest of) the result to a pandas DataFrame.

This method is only available if the `pandas` library is installed.
Expand All @@ -540,7 +544,7 @@ async def to_df(self, expand=False):
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
:param expand: If :const:`True`, some structures in the result will be
recursively expanded (flattened out into multiple columns) like so
(everything inside ``<...>`` is a placeholder):

Expand Down Expand Up @@ -604,6 +608,11 @@ async def to_df(self, expand=False):
:const:`dict` keys and variable names that contain ``.`` or ``\``
will be escaped with a backslash (``\.`` and ``\\`` respectively).
:type expand: bool
:param parse_dates:
If :const:`True`, columns that excluvively contain
:class:`time.DateTime` objects, :class:`time.Date` objects, or
:const:`None`, will be converted to :class:`pandas.Timestamp`.
:type parse_dates: bool

:rtype: :py:class:`pandas.DataFrame`
:raises ImportError: if `pandas` library is not available.
Expand All @@ -618,7 +627,7 @@ async def to_df(self, expand=False):
import pandas as pd

if not expand:
return pd.DataFrame(await self.values(), columns=self._keys)
df = pd.DataFrame(await self.values(), columns=self._keys)
else:
df_keys = None
rows = []
Expand All @@ -638,13 +647,29 @@ async def to_df(self, expand=False):
df_keys = False
rows.append(row)
if df_keys is False:
return pd.DataFrame(rows)
df = pd.DataFrame(rows)
else:
columns = df_keys or [
k.replace(".", "\\.").replace("\\", "\\\\")
for k in self._keys
]
return pd.DataFrame(rows, columns=columns)
df = pd.DataFrame(rows, columns=columns)
if not parse_dates:
return df
dt_columns = df.columns[df.apply(
lambda col: pd.api.types.infer_dtype(col) == "mixed" and col.map(
lambda x: isinstance(x, (DateTime, Date, type(None)))
).all()
)]
df[dt_columns] = df[dt_columns].apply(
lambda col: col.map(
lambda x:
pd.Timestamp(x.iso_format())
.replace(tzinfo=getattr(x, "tzinfo", None))
if x else pd.NaT
)
)
return df

def closed(self):
"""Return True if the result has been closed.
Expand Down
35 changes: 30 additions & 5 deletions neo4j/_sync/work/result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@
ResultNotSingleError,
)
from ...meta import experimental
from ...time import (
Date,
DateTime,
)
from ...work import ResultSummary
from ..io import ConnectionErrorHandler

Expand Down Expand Up @@ -527,7 +531,7 @@ def data(self, *keys):

@experimental("pandas support is experimental and might be changed or "
"removed in future versions")
def to_df(self, expand=False):
def to_df(self, expand=False, parse_dates=False):
r"""Convert (the rest of) the result to a pandas DataFrame.

This method is only available if the `pandas` library is installed.
Expand All @@ -540,7 +544,7 @@ def to_df(self, expand=False):
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
:param expand: If :const:`True`, some structures in the result will be
recursively expanded (flattened out into multiple columns) like so
(everything inside ``<...>`` is a placeholder):

Expand Down Expand Up @@ -604,6 +608,11 @@ def to_df(self, expand=False):
:const:`dict` keys and variable names that contain ``.`` or ``\``
will be escaped with a backslash (``\.`` and ``\\`` respectively).
:type expand: bool
:param parse_dates:
If :const:`True`, columns that excluvively contain
:class:`time.DateTime` objects, :class:`time.Date` objects, or
:const:`None`, will be converted to :class:`pandas.Timestamp`.
:type parse_dates: bool

:rtype: :py:class:`pandas.DataFrame`
:raises ImportError: if `pandas` library is not available.
Expand All @@ -618,7 +627,7 @@ def to_df(self, expand=False):
import pandas as pd

if not expand:
return pd.DataFrame(self.values(), columns=self._keys)
df = pd.DataFrame(self.values(), columns=self._keys)
else:
df_keys = None
rows = []
Expand All @@ -638,13 +647,29 @@ def to_df(self, expand=False):
df_keys = False
rows.append(row)
if df_keys is False:
return pd.DataFrame(rows)
df = pd.DataFrame(rows)
else:
columns = df_keys or [
k.replace(".", "\\.").replace("\\", "\\\\")
for k in self._keys
]
return pd.DataFrame(rows, columns=columns)
df = pd.DataFrame(rows, columns=columns)
if not parse_dates:
return df
dt_columns = df.columns[df.apply(
lambda col: pd.api.types.infer_dtype(col) == "mixed" and col.map(
lambda x: isinstance(x, (DateTime, Date, type(None)))
).all()
)]
df[dt_columns] = df[dt_columns].apply(
lambda col: col.map(
lambda x:
pd.Timestamp(x.iso_format())
.replace(tzinfo=getattr(x, "tzinfo", None))
if x else pd.NaT
)
)
return df

def closed(self):
"""Return True if the result has been closed.
Expand Down
4 changes: 4 additions & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@
install_requires = [
"pytz",
]
extra_require = {
"pandas": ["pandas>=1.0.0"],
}
classifiers = [
"Intended Audience :: Developers",
"License :: OSI Approved :: Apache Software License",
Expand Down Expand Up @@ -67,6 +70,7 @@
"keywords": "neo4j graph database",
"url": "https://github.com/neo4j/neo4j-python-driver",
"install_requires": install_requires,
"extra_require": extra_require,
"classifiers": classifiers,
"packages": packages,
"entry_points": entry_points,
Expand Down
Loading