Accept more types as query parameters (#881)

Implement serialization support for
 * python types: tuple
 * pandas types:
   * all but period, interval, pyarrow
 * numpy types:
   * all but void, complexfloating

pyarrow support was not implemented as it would either require more ifs in the
recursive packing function, making it (driver's hot-path) slower for non-pyarrow
use-cases. Or alternatively, transformers would have to be used making pyarrow
type serialization rather slow.

Co-authored-by: Florent Biville <[email protected]>
Co-authored-by: Grant Lodge <[email protected]>

Author: 3 people

docs/source/_images/core_type_mappings.png

@@ -1099,7 +1099,7 @@ The core types with their general mappings are listed below:
 +------------------------+---------------------------------------------------------------------------------------------------------------------------+
 | String                 | :class:`str`                                                                                                              |
 +------------------------+---------------------------------------------------------------------------------------------------------------------------+
-| Bytes :sup:`[1]`       | :class:`bytearray`                                                                                                        |
+| Bytes :sup:`[1]`       | :class:`bytes`                                                                                                            |
 +------------------------+---------------------------------------------------------------------------------------------------------------------------+
 | List                   | :class:`list`                                                                                                             |
 +------------------------+---------------------------------------------------------------------------------------------------------------------------+
@@ -1118,6 +1118,57 @@ The diagram below illustrates the actual mappings between the various layers, fr
     :target: ./_images/core_type_mappings.svg
 
 
+Extended Data Types
+===================
+
+The driver supports serializing more types (as parameters in).
+However, they will have to be mapped to the existing Bolt types (see above) when they are sent to the server.
+This means, the driver will never return these types in results.
+
+When in doubt, you can test the type conversion like so::
+
+    import neo4j
+
+
+    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
+        with driver.session() as session:
+            type_in = ("foo", "bar")
+            result = session.run("RETURN $x", x=type_in)
+            type_out = result.single()[0]
+            print(type(type_out))
+            print(type_out)
+
+Which in this case would yield::
+
+    <class 'list'>
+    ['foo', 'bar']
+
+
++-----------------------------------+---------------------------------+---------------------------------------+
+| Parameter Type                    | Bolt Type                       | Result Type                           |
++===================================+=================================+=======================================+
+| :class:`tuple`                    | List                            | :class:`list`                         |
++-----------------------------------+---------------------------------+---------------------------------------+
+| :class:`bytearray`                | Bytes                           | :class:`bytes`                        |
++-----------------------------------+---------------------------------+---------------------------------------+
+| numpy\ :sup:`[2]` ``ndarray``     | (nested) List                   | (nested) :class:`list`                |
++-----------------------------------+---------------------------------+---------------------------------------+
+| pandas\ :sup:`[3]` ``DataFrame``  | Map[str, List[_]] :sup:`[4]`    | :class:`dict`                         |
++-----------------------------------+---------------------------------+---------------------------------------+
+| pandas ``Series``                 | List                            | :class:`list`                         |
++-----------------------------------+---------------------------------+---------------------------------------+
+| pandas ``Array``                  | List                            | :class:`list`                         |
++-----------------------------------+---------------------------------+---------------------------------------+
+
+.. Note::
+
+   2. ``void`` and ``complexfloating`` typed numpy ``ndarray``\s are not supported.
+   3. ``Period``, ``Interval``, and ``pyarrow`` pandas types are not supported.
+   4. A pandas ``DataFrame`` will be serialized as Map with the column names mapping to the column values (as Lists).
+       Just like with ``dict`` objects, the column names need to be :class:`str` objects.
+
+
+
 ****************
 Graph Data Types
 ****************

docs/source/_images/core_type_mappings.svg

@@ -45,10 +45,17 @@ dynamic = ["version", "readme"]
 Homepage = "https://github.com/neo4j/neo4j-python-driver"
 
 [project.optional-dependencies]
-pandas = ["pandas>=1.0.0"]
+numpy = ["numpy >= 1.7.0, < 2.0.0"]
+pandas = [
+    "pandas >= 1.1.0, < 2.0.0",
+    "numpy >= 1.7.0, < 2.0.0",
+]
 
 [build-system]
-requires = ["setuptools~=65.6", "tomlkit~=0.11.6"]
+requires = [
+    "setuptools~=65.6",
+    "tomlkit~=0.11.6",
+]
 build-backend = "setuptools.build_meta"
 
 # still in beta

docs/source/api.rst

@@ -18,7 +18,9 @@ tomlkit~=0.11.6
 # needed for running tests
 coverage[toml]>=5.5
 mock>=4.0.3
+numpy>=1.7.0
 pandas>=1.0.0
+pyarrow>=1.0.0
 pytest>=6.2.5
 pytest-asyncio>=0.16.0
 pytest-benchmark>=3.4.1

pyproject.toml

@@ -17,13 +17,15 @@
 
 from ._common import (
     BrokenHydrationObject,
+    DehydrationHooks,
     HydrationScope,
 )
 from ._interface import HydrationHandlerABC
 
 
 __all__ = [
     "BrokenHydrationObject",
+    "DehydrationHooks",
     "HydrationHandlerABC",
     "HydrationScope",
 ]

requirements-dev.txt

@@ -16,12 +16,51 @@
 # limitations under the License.
 
 
+import typing as t
 from copy import copy
+from dataclasses import dataclass
 
 from ...graph import Graph
 from ..packstream import Structure
 
 
+@dataclass
+class DehydrationHooks:
+    exact_types: t.Dict[t.Type, t.Callable[[t.Any], t.Any]]
+    subtypes: t.Dict[t.Type, t.Callable[[t.Any], t.Any]]
+
+    def update(self, exact_types=None, subtypes=None):
+        exact_types = exact_types or {}
+        subtypes = subtypes or {}
+        self.exact_types.update(exact_types)
+        self.subtypes.update(subtypes)
+
+    def extend(self, exact_types=None, subtypes=None):
+        exact_types = exact_types or {}
+        subtypes = subtypes or {}
+        return DehydrationHooks(
+            exact_types={**self.exact_types, **exact_types},
+            subtypes={**self.subtypes, **subtypes},
+        )
+
+    def get_transformer(self, item):
+        type_ = type(item)
+        transformer = self.exact_types.get(type_)
+        if transformer is not None:
+            return transformer
+        transformer = next(
+            (
+                f
+                for super_type, f in self.subtypes.items()
+                if isinstance(item, super_type)
+            ),
+            None,
+        )
+        if transformer is not None:
+            return transformer
+        return None
+
+
 class BrokenHydrationObject:
     """
     Represents an object from the server, not understood by the driver.
@@ -68,7 +107,7 @@ def __init__(self, hydration_handler, graph_hydrator):
             list: self._hydrate_list,
             dict: self._hydrate_dict,
         }
-        self.dehydration_hooks = hydration_handler.dehydration_functions
+        self.dehydration_hooks = hydration_handler.dehydration_hooks
 
     def _hydrate_structure(self, value):
         f = self._struct_hydration_functions.get(value.tag)

src/neo4j/_codec/hydration/__init__.py

@@ -18,11 +18,14 @@
 
 import abc
 
+from .._common import DehydrationHooks
+
 
 class HydrationHandlerABC(abc.ABC):
     def __init__(self):
         self.struct_hydration_functions = {}
-        self.dehydration_functions = {}
+        self.dehydration_hooks = DehydrationHooks(exact_types={},
+                                                  subtypes={})
 
     @abc.abstractmethod
     def new_hydration_scope(self):

src/neo4j/_codec/hydration/_common.py

@@ -23,6 +23,10 @@
     timedelta,
 )
 
+from ...._optional_deps import (
+    np,
+    pd,
+)
 from ....graph import (
     Graph,
     Node,
@@ -159,8 +163,7 @@ def __init__(self):
             b"d": temporal.hydrate_datetime,     # no time zone
             b"E": temporal.hydrate_duration,
         }
-        self.dehydration_functions = {
-            **self.dehydration_functions,
+        self.dehydration_hooks.update(exact_types={
             Point: spatial.dehydrate_point,
             CartesianPoint: spatial.dehydrate_point,
             WGS84Point: spatial.dehydrate_point,
@@ -172,7 +175,19 @@ def __init__(self):
             datetime: temporal.dehydrate_datetime,
             Duration: temporal.dehydrate_duration,
             timedelta: temporal.dehydrate_timedelta,
-        }
+        })
+        if np is not None:
+            self.dehydration_hooks.update(exact_types={
+                np.datetime64: temporal.dehydrate_np_datetime,
+                np.timedelta64: temporal.dehydrate_np_timedelta,
+            })
+        if pd is not None:
+            self.dehydration_hooks.update(exact_types={
+                pd.Timestamp: temporal.dehydrate_pandas_datetime,
+                pd.Timedelta: temporal.dehydrate_pandas_timedelta,
+                type(pd.NaT): lambda _: None,
+            })
+
 
     def patch_utc(self):
         from ..v2 import temporal as temporal_v2
@@ -186,10 +201,18 @@ def patch_utc(self):
             b"i": temporal_v2.hydrate_datetime,
         })
 
-        self.dehydration_functions.update({
+        self.dehydration_hooks.update(exact_types={
             DateTime: temporal_v2.dehydrate_datetime,
             datetime: temporal_v2.dehydrate_datetime,
         })
+        if np is not None:
+            self.dehydration_hooks.update(exact_types={
+                np.datetime64: temporal_v2.dehydrate_np_datetime,
+            })
+        if pd is not None:
+            self.dehydration_hooks.update(exact_types={
+                pd.Timestamp: temporal_v2.dehydrate_pandas_datetime,
+            })
 
     def new_hydration_scope(self):
         self._created_scope = True

src/neo4j/_codec/hydration/_interface/__init__.py

@@ -22,10 +22,17 @@
     timedelta,
 )
 
+from ...._optional_deps import (
+    np,
+    pd,
+)
 from ....time import (
     Date,
     DateTime,
     Duration,
+    MAX_YEAR,
+    MIN_YEAR,
+    NANO_SECONDS,
     Time,
 )
 from ...packstream import Structure
@@ -171,6 +178,50 @@ def seconds_and_nanoseconds(dt):
                          int(tz.utcoffset(value).total_seconds()))
 
 
+if np is not None:
+    def dehydrate_np_datetime(value):
+        """ Dehydrator for `numpy.datetime64` values.
+
+        :param value:
+        :type value: numpy.datetime64
+        :returns:
+        """
+        if np.isnat(value):
+            return None
+        year = value.astype("datetime64[Y]").astype(int) + 1970
+        if not 0 < year <= 9999:
+            # while we could encode years outside the range, they would fail
+            # when retrieved from the database.
+            raise ValueError(f"Year out of range ({MIN_YEAR:d}..{MAX_YEAR:d}) "
+                             f"found {year}")
+        seconds = value.astype(np.dtype("datetime64[s]")).astype(int)
+        nanoseconds = (value.astype(np.dtype("datetime64[ns]")).astype(int)
+                       % NANO_SECONDS)
+        return Structure(b"d", seconds, nanoseconds)
+
+
+if pd is not None:
+    def dehydrate_pandas_datetime(value):
+        """ Dehydrator for `pandas.Timestamp` values.
+
+        :param value:
+        :type value: pandas.Timestamp
+        :returns:
+        """
+        return dehydrate_datetime(
+            DateTime(
+                value.year,
+                value.month,
+                value.day,
+                value.hour,
+                value.minute,
+                value.second,
+                value.microsecond * 1000 + value.nanosecond,
+                value.tzinfo,
+            )
+        )
+
+
 def hydrate_duration(months, days, seconds, nanoseconds):
     """ Hydrator for `Duration` values.
 
@@ -205,3 +256,50 @@ def dehydrate_timedelta(value):
     seconds = value.seconds
     nanoseconds = 1000 * value.microseconds
     return Structure(b"E", months, days, seconds, nanoseconds)
+
+
+if np is not None:
+    _NUMPY_DURATION_UNITS = {
+        "Y": "years",
+        "M": "months",
+        "W": "weeks",
+        "D": "days",
+        "h": "hours",
+        "m": "minutes",
+        "s": "seconds",
+        "ms": "milliseconds",
+        "us": "microseconds",
+        "ns": "nanoseconds",
+    }
+
+    def dehydrate_np_timedelta(value):
+        """ Dehydrator for `numpy.timedelta64` values.
+
+        :param value:
+        :type value: numpy.timedelta64
+        :returns:
+        """
+        if np.isnat(value):
+            return None
+        unit, step_size = np.datetime_data(value)
+        numer = int(value.astype(int))
+        # raise RuntimeError((type(numer), type(step_size)))
+        kwarg = _NUMPY_DURATION_UNITS.get(unit)
+        if kwarg is not None:
+            return dehydrate_duration(Duration(**{kwarg: numer * step_size}))
+        return dehydrate_duration(Duration(
+            nanoseconds=value.astype("timedelta64[ns]").astype(int)
+        ))
+
+
+if pd is not None:
+    def dehydrate_pandas_timedelta(value):
+        """ Dehydrator for `pandas.Timedelta` values.
+
+        :param value:
+        :type value: pandas.Timedelta
+        :returns:
+        """
+        return dehydrate_duration(Duration(
+            nanoseconds=value.value
+        ))