[pymeos integration] Part 2 - serialization using pymeos

adonmo · Jun 14, 2020 · 5cab3e0 · 5cab3e0
1 parent c7a1902
commit 5cab3e0
Show file tree

Hide file tree

Showing 9 changed files with 70 additions and 48 deletions.
diff --git a/mobilitydb_sqlalchemy/types/TBaseType.py b/mobilitydb_sqlalchemy/types/TBaseType.py
@@ -19,17 +19,29 @@ class TBaseType(UserDefinedType):
     pandas_value_column = "value"
     comparator_factory = Comparator
 
+    @property
+    def pymeos_sequence_type(self):
+        raise NotImplementedError()
+
+    @property
+    def pymeos_instant_type(self):
+        raise NotImplementedError()
+
     @property
     def pymeos_deserializer_type(self):
         raise NotImplementedError()
 
+    @property
+    def pymeos_serializer_type(self):
+        raise NotImplementedError()
+
     @staticmethod
     def validate_type(value):
         raise NotImplementedError()
 
     @staticmethod
     def write_instant_value(value):
-        raise NotImplementedError()
+        return value
 
     @staticmethod
     def parse_instant_value(value):
@@ -43,20 +55,21 @@ def bind_processor(self, dialect):
         def process(value):
             self.validate_type(value)
             value.sort_index(inplace=True)
-            instants = list(value.loc[v] for v in value.index)
-            left_bound = "[" if self.left_closed else "("
-            right_bound = "]" if self.right_closed else ")"
-            return "{}{}{}".format(
-                left_bound,
-                ", ".join(
-                    "{}@{}".format(
-                        self.write_instant_value(getattr(i, self.pandas_value_column)),
-                        i.name,
-                    )
-                    for i in instants
-                ),
-                right_bound,
+
+            serializer = self.pymeos_serializer_type()
+            instants = list(
+                self.pymeos_instant_type(
+                    self.write_instant_value(
+                        getattr(value.loc[v], self.pandas_value_column)
+                    ),
+                    int(v.timestamp() * 1000),
+                )
+                for v in value.index
+            )
+            sequence = self.pymeos_sequence_type(
+                instants, not self.left_closed, not self.right_closed
             )
+            return serializer.write(sequence)
 
         return process
 

diff --git a/mobilitydb_sqlalchemy/types/TBool.py b/mobilitydb_sqlalchemy/types/TBool.py
@@ -1,13 +1,16 @@
 import pandas as pd
 
-from pymeos import DeserializerBool
+from pymeos import DeserializerBool, SerializerBool, TInstantBool, TSequenceBool
 from sqlalchemy.types import UserDefinedType
 
 from .TBaseType import TBaseType
 
 
 class TBool(TBaseType):
+    pymeos_sequence_type = TSequenceBool
+    pymeos_instant_type = TInstantBool
     pymeos_deserializer_type = DeserializerBool
+    pymeos_serializer_type = SerializerBool
 
     def get_col_spec(self):
         return "TBOOL"
@@ -19,7 +22,3 @@ def validate_type(value):
             raise TypeError(
                 "TBool needs bool values. Got: {} Expected {}".format(dtype_kind, "b")
             )
-
-    @staticmethod
-    def write_instant_value(py_bool):
-        return "t" if py_bool else "f"
diff --git a/mobilitydb_sqlalchemy/types/TFloat.py b/mobilitydb_sqlalchemy/types/TFloat.py
@@ -1,14 +1,17 @@
 import pandas as pd
 
-from pymeos import DeserializerFloat
+from pymeos import DeserializerFloat, SerializerFloat, TInstantFloat, TSequenceFloat
 from sqlalchemy.types import UserDefinedType
 from pandas.api.types import is_numeric_dtype
 
 from .TBaseType import TBaseType
 
 
 class TFloat(TBaseType):
+    pymeos_sequence_type = TSequenceFloat
+    pymeos_instant_type = TInstantFloat
     pymeos_deserializer_type = DeserializerFloat
+    pymeos_serializer_type = SerializerFloat
 
     def get_col_spec(self):
         return "TFLOAT"
@@ -20,7 +23,3 @@ def validate_type(value):
             raise TypeError(
                 "TInt needs int values. Got: {} Expected {}".format(dtypes.kind, "f")
             )
-
-    @staticmethod
-    def write_instant_value(value):
-        return str(value)
diff --git a/mobilitydb_sqlalchemy/types/TGeogPoint.py b/mobilitydb_sqlalchemy/types/TGeogPoint.py
@@ -1,6 +1,12 @@
 import re
 
-from pymeos import DeserializerGeom
+from pymeos import (
+    DeserializerGeom,
+    Geometry,
+    SerializerGeom,
+    TInstantGeom,
+    TSequenceGeom,
+)
 from sqlalchemy import func
 from sqlalchemy.types import UserDefinedType
 
@@ -24,7 +30,10 @@
 
 
 class TGeogPoint(TBaseType):
+    pymeos_sequence_type = TSequenceGeom
+    pymeos_instant_type = TInstantGeom
     pymeos_deserializer_type = DeserializerGeom
+    pymeos_serializer_type = SerializerGeom
 
     # This is ensure compatibility with movingpandas
     pandas_value_column = "geometry"
@@ -47,7 +56,7 @@ def validate_type(value):
 
     @staticmethod
     def write_instant_value(value):
-        return value.wkt
+        return Geometry(value.wkt)
 
     @staticmethod
     def parse_instant_value(value):

diff --git a/mobilitydb_sqlalchemy/types/TGeomPoint.py b/mobilitydb_sqlalchemy/types/TGeomPoint.py
@@ -1,6 +1,12 @@
 import re
 
-from pymeos import DeserializerGeom
+from pymeos import (
+    DeserializerGeom,
+    Geometry,
+    SerializerGeom,
+    TInstantGeom,
+    TSequenceGeom,
+)
 from sqlalchemy import func
 from sqlalchemy.types import UserDefinedType
 
@@ -24,7 +30,10 @@
 
 
 class TGeomPoint(TBaseType):
+    pymeos_sequence_type = TSequenceGeom
+    pymeos_instant_type = TInstantGeom
     pymeos_deserializer_type = DeserializerGeom
+    pymeos_serializer_type = SerializerGeom
 
     # This is ensure compatibility with movingpandas
     pandas_value_column = "geometry"
@@ -47,7 +56,7 @@ def validate_type(value):
 
     @staticmethod
     def write_instant_value(value):
-        return value.wkt
+        return Geometry(value.wkt)
 
     @staticmethod
     def parse_instant_value(value):

diff --git a/mobilitydb_sqlalchemy/types/TInt.py b/mobilitydb_sqlalchemy/types/TInt.py
@@ -22,7 +22,3 @@ def validate_type(value):
             raise TypeError(
                 "TInt needs int values. Got: {} Expected {}".format(dtype_kind, "i")
             )
-
-    @staticmethod
-    def write_instant_value(value):
-        return str(value)
diff --git a/mobilitydb_sqlalchemy/utils.py b/mobilitydb_sqlalchemy/utils.py
@@ -0,0 +1,11 @@
+import datetime
+import pytz
+
+
+def epoch(year, month, day, hour=0, minute=0):
+    return int(
+        datetime.datetime(year, month, day, hour, minute)
+        .replace(tzinfo=pytz.UTC)
+        .timestamp()
+        * 1000
+    )
diff --git a/tests/test_tgeogpoint.py b/tests/test_tgeogpoint.py
@@ -12,6 +12,7 @@
 from sqlalchemy.exc import StatementError
 from sqlalchemy.sql.expression import cast
 
+from mobilitydb_sqlalchemy.utils import epoch
 from .models import GeogTrips, GeogTripsWithMovingPandas
 from .postgis_types import Geometry
 
@@ -216,14 +217,6 @@ def test_mobility_functions(session):
         .all()
     )
 
-    def epoch(year, month, day, hour=0, minute=0):
-        return int(
-            datetime.datetime(year, month, day, hour, minute)
-            .replace(tzinfo=pytz.UTC)
-            .timestamp()
-            * 1000
-        )
-
     assert len(trips) == 1
     assert trips[0][0] == 10
     assert trips[0][1] == 20

diff --git a/tests/test_tgeompoint.py b/tests/test_tgeompoint.py
@@ -10,6 +10,7 @@
 from sqlalchemy import alias, func
 from sqlalchemy.exc import StatementError
 
+from mobilitydb_sqlalchemy.utils import epoch
 from .models import Trips, TripsWithMovingPandas
 
 
@@ -203,14 +204,6 @@ def test_mobility_functions(session):
         .all()
     )
 
-    def epoch(year, month, day, hour=0, minute=0):
-        return int(
-            datetime.datetime(year, month, day, hour, minute)
-            .replace(tzinfo=pytz.UTC)
-            .timestamp()
-            * 1000
-        )
-
     assert len(trips) == 1
     assert trips[0][0] == 10
     assert trips[0][1] == 20