Added new geo functions to omniscidb backend

ibis/expr/api.py

@@ -678,12 +678,21 @@ def cast(arg, target_type):
     if op.to.equals(arg.type()):
         # noop case if passed type is the same
         return arg
-    else:
-        result = op.to_expr()
-        if not arg.has_name():
-            return result
-        expr_name = 'cast({}, {})'.format(arg.get_name(), op.to)
-        return result.name(expr_name)
+
+    if isinstance(op.to, (dt.Geography, dt.Geometry)):
+        from_geotype = arg.type().geotype
+        to_geotype = op.to.geotype
+        if (
+            from_geotype == to_geotype
+            or (from_geotype is None and to_geotype == 'geometry')
+        ):
+            return arg
+
+    result = op.to_expr()
+    if not arg.has_name():
+        return result
+    expr_name = 'cast({}, {})'.format(arg.get_name(), op.to)
+    return result.name(expr_name)
 
 
 cast.__doc__ = """
@@ -2148,6 +2157,22 @@ def geo_srid(arg):
     return op.to_expr()
 
 
+def geo_set_srid(arg, srid):
+    """Set the spatial reference identifier for the ST_Geometry
+
+    Parameters
+    ----------
+    arg : geometry
+    srid : integer
+
+    Returns
+    -------
+    SetSRID : geometry
+    """
+    op = ops.GeoSetSRID(arg, srid)
+    return op.to_expr()
+
+
 def geo_buffer(arg, radius):
     """Returns a geometry that represents all points whose distance from this
     Geometry is less than or equal to distance. Calculations are in the
@@ -2331,6 +2356,7 @@ def geo_transform(arg, srid):
     overlaps=geo_overlaps,
     perimeter=geo_perimeter,
     point_n=geo_point_n,
+    set_srid=geo_set_srid,
     simplify=geo_simplify,
     srid=geo_srid,
     start_point=geo_start_point,

ibis/expr/datatypes.py

@@ -613,6 +613,46 @@ def __init__(
         self.geotype = geotype
         self.srid = srid
 
+    def __str__(self) -> str:
+        geo_op = self.name.lower()
+        if self.geotype is not None:
+            geo_op += ':' + self.geotype
+        if self.srid is not None:
+            geo_op += ';' + str(self.srid)
+        return geo_op
+
+
+class Geometry(GeoSpatial):
+    """Geometry is used to cast from geography types."""
+
+    column = ir.GeoSpatialColumn
+    scalar = ir.GeoSpatialScalar
+
+    __slots__ = ()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.geotype = 'geometry'
+
+    def __str__(self) -> str:
+        return self.name.lower()
+
+
+class Geography(GeoSpatial):
+    """Geography is used to cast from geometry types."""
+
+    column = ir.GeoSpatialColumn
+    scalar = ir.GeoSpatialScalar
+
+    __slots__ = ()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.geotype = 'geography'
+
+    def __str__(self) -> str:
+        return self.name.lower()
+
 
 class Point(GeoSpatial):
     """A point described by two coordinates."""
@@ -1041,6 +1081,10 @@ def type(self) -> DataType:
 
         field : [a-zA-Z_][a-zA-Z_0-9]*
 
+        geography: "geography"
+
+        geometry: "geometry"
+
         point : "point"
               | "point" ";" srid
               | "point" ":" geotype
@@ -1176,6 +1220,12 @@ def type(self) -> DataType:
             return Struct(names, types)
 
         # geo spatial data type
+        elif self._accept(Tokens.GEOMETRY):
+            return Geometry()
+
+        elif self._accept(Tokens.GEOGRAPHY):
+            return Geography()
+
         elif self._accept(Tokens.POINT):
             geotype = None
             srid = None
@@ -1516,10 +1566,10 @@ def can_cast_variadic(
 
 
 # geo spatial data type
-@castable.register(Array, Point)
-@castable.register(Array, LineString)
-@castable.register(Array, Polygon)
-@castable.register(Array, MultiPolygon)
+# cast between same type, used to cast from/to geometry and geography
+@castable.register(Array, (Point, LineString, Polygon, MultiPolygon))
+@castable.register((Point, LineString, Polygon, MultiPolygon), Geometry)
+@castable.register((Point, LineString, Polygon, MultiPolygon), Geography)
 def can_cast_geospatial(source, target, **kwargs):
     return True
 

ibis/expr/operations.py

@@ -3176,11 +3176,17 @@ class GeoNRings(GeoSpatialUnOp):
 
 
 class GeoSRID(GeoSpatialUnOp):
-    """Returns the spatial reference identifier for the ST_Geometry"""
+    """Returns the spatial reference identifier for the ST_Geometry."""
 
     output_type = rlz.shape_like('args', dt.int64)
 
 
+class GeoSetSRID(GeoSpatialUnOp):
+    """Set the spatial reference identifier for the ST_Geometry."""
+    srid = Arg(rlz.integer)
+    output_type = rlz.shape_like('args', dt.geometry)
+
+
 class GeoBuffer(GeoSpatialUnOp):
     """Returns a geometry that represents all points whose distance from this
     Geometry is less than or equal to distance. Calculations are in the

ibis/expr/tests/test_geospatial.py

@@ -45,6 +45,7 @@ def test_geo_ops_smoke(geo_table):
 
     # test ops
     point.srid()
+    point.set_srid(4326)
     point.x()
     point.y()
 

ibis/mapd/operations.py

@@ -54,8 +54,18 @@ def _cast(translator, expr):
     op = expr.op()
     arg, target = op.args
     arg_ = translator.translate(arg)
-    type_ = str(MapDDataType.from_ibis(target, nullable=False))
 
+    if isinstance(arg, ir.GeoSpatialValue):
+        # NOTE: CastToGeography expects geometry with SRID=4326
+        type_ = target.geotype.upper()
+
+        if type_ == 'GEOMETRY':
+            raise com.UnsupportedOperationError(
+                'OmnisciDB/MapD doesn\'t support yet convert '
+                + 'from GEOGRAPHY to GEOMETRY.'
+            )
+    else:
+        type_ = str(MapDDataType.from_ibis(target, nullable=False))
     return 'CAST({0!s} AS {1!s})'.format(arg_, type_)
 
 
@@ -812,6 +822,8 @@ def formatter(translator, expr):
     ops.GeoNPoints: unary('ST_NPOINTS'),
     ops.GeoNRings: unary('ST_NRINGS'),
     ops.GeoSRID: unary('ST_SRID'),
+    ops.GeoTransform: fixed_arity('ST_TRANSFORM', 2),
+    ops.GeoSetSRID: fixed_arity('ST_SETSRID', 2)
 }
 
 # STRING

ibis/tests/all/test_geospatial.py

@@ -1,6 +1,7 @@
 """ Tests for geo spatial data types"""
 from inspect import isfunction
 
+import numpy as np
 import pytest
 from numpy import testing
 
@@ -11,6 +12,8 @@
 point_0 = ibis.literal((0, 0), type='point:geometry').name('p')
 point_1 = ibis.literal((1, 1), type='point:geometry').name('p')
 point_2 = ibis.literal((2, 2), type='point;4326:geometry').name('p')
+point_3 = ibis.literal((1, 1), type='point:geography').name('p')
+point_4 = ibis.literal((2, 2), type='point;4326:geography').name('p')
 polygon_0 = ibis.literal(
     (
         ((1, 0), (0, 1), (-1, 0), (0, -1), (1, 0)),
@@ -19,6 +22,9 @@
     type='polygon',
 )
 
+# add here backends that support geo spatial types
+all_db_geo_supported = [MapD]
+
 
 @pytest.mark.parametrize(
     ('expr_fn', 'expected'),
@@ -36,7 +42,7 @@
         (lambda t: t['geo_point'].srid(), [0] * 5),
     ],
 )
-@pytest.mark.only_on_backends([MapD])
+@pytest.mark.only_on_backends(all_db_geo_supported)
 def test_geo_spatial_unops(backend, geo, expr_fn, expected):
     """Testing for geo spatial unary operations."""
     expr = expr_fn(geo)
@@ -67,7 +73,7 @@ def test_geo_spatial_unops(backend, geo, expr_fn, expected):
         ),
     ],
 )
-@pytest.mark.only_on_backends([MapD])
+@pytest.mark.only_on_backends(all_db_geo_supported)
 def test_geo_spatial_binops(backend, geo, fn, arg_left, arg_right, expected):
     """Testing for geo spatial binary operations."""
     left = arg_left(geo) if isfunction(arg_left) else arg_left
@@ -85,7 +91,7 @@ def test_geo_spatial_binops(backend, geo, fn, arg_left, arg_right, expected):
         (lambda t: t['geo_linestring'].start_point(), [False] * 5),
     ],
 )
-@pytest.mark.only_on_backends([MapD])
+@pytest.mark.only_on_backends(all_db_geo_supported)
 def test_get_point(backend, geo, expr_fn, expected):
     """Testing for geo spatial get point operations."""
     # a geo spatial data does not contain its boundary
@@ -96,17 +102,117 @@ def test_get_point(backend, geo, expr_fn, expected):
 
 
 @pytest.mark.parametrize(('arg', 'expected'), [(polygon_0, [1.98] * 5)])
-@pytest.mark.only_on_backends([MapD])
+@pytest.mark.only_on_backends(all_db_geo_supported)
 def test_area(backend, geo, arg, expected):
     """Testing for geo spatial area operation."""
-    expr = geo[geo, arg.area().name('tmp')]
+    expr = geo[geo.id, arg.area().name('tmp')]
     result = expr.execute()['tmp']
     testing.assert_almost_equal(result, expected, decimal=2)
 
 
-@pytest.mark.parametrize(('arg', 'expected'), [(point_2.srid() == 4326, True)])
-@pytest.mark.only_on_backends([MapD])
-def test_srid_literals(backend, geo, arg, expected):
-    """Testing for geo spatial srid operation (from literal)."""
-    result = geo[geo, arg.name('tmp')].execute()['tmp'][[0]]
-    testing.assert_almost_equal(result, [expected], decimal=2)
+@pytest.mark.parametrize(
+    ('condition', 'expected'),
+    [
+        (lambda t: point_2.srid(), 4326),
+        (lambda t: point_0.srid(), 0),
+        (lambda t: t.geo_point.srid(), 0),
+        (lambda t: t.geo_linestring.srid(), 0),
+        (lambda t: t.geo_polygon.srid(), 0),
+        (lambda t: t.geo_multipolygon.srid(), 0),
+    ]
+)
+@pytest.mark.only_on_backends(all_db_geo_supported)
+def test_srid(backend, geo, condition, expected):
+    """Testing for geo spatial srid operation."""
+    expr = geo[geo.id, condition(geo).name('tmp')]
+    result = expr.execute()['tmp'][[0]]
+    assert np.all(result == expected)
+
+
+@pytest.mark.parametrize(
+    ('condition', 'expected'),
+    [
+        (lambda t: point_0.set_srid(4326).srid(), 4326),
+        (lambda t: point_0.set_srid(4326).set_srid(0).srid(), 0),
+        (lambda t: t.geo_point.set_srid(4326).srid(), 4326),
+        (lambda t: t.geo_linestring.set_srid(4326).srid(), 4326),
+        (lambda t: t.geo_polygon.set_srid(4326).srid(), 4326),
+        (lambda t: t.geo_multipolygon.set_srid(4326).srid(), 4326),
+    ]
+)
+@pytest.mark.only_on_backends(all_db_geo_supported)
+def test_set_srid(backend, geo, condition, expected):
+    """Testing for geo spatial set_srid operation."""
+    expr = geo[geo.id, condition(geo).name('tmp')]
+    result = expr.execute()['tmp'][[0]]
+    assert np.all(result == expected)
+
+
+@pytest.mark.parametrize(
+    ('condition', 'expected'),
+    [
+        (lambda t: point_0.set_srid(4326).transform(900913).srid(),
+         900913),
+        (lambda t: point_2.transform(900913).srid(),
+         900913),
+        (lambda t: t.geo_point.set_srid(4326).transform(900913).srid(),
+         900913),
+        (lambda t: t.geo_linestring.set_srid(4326).transform(900913).srid(),
+         900913),
+        (lambda t: t.geo_polygon.set_srid(4326).transform(900913).srid(),
+         900913),
+        (lambda t: t.geo_multipolygon.set_srid(4326).transform(900913).srid(),
+         900913),
+    ]
+)
+@pytest.mark.only_on_backends(all_db_geo_supported)
+@pytest.mark.xfail_unsupported
+def test_transform(backend, geo, condition, expected):
+    """Testing for geo spatial transform operation."""
+    expr = geo[geo.id, condition(geo).name('tmp')]
+    result = expr.execute()['tmp'][[0]]
+    assert np.all(result == expected)
+
+
+@pytest.mark.parametrize(
+    'expr_fn',
+    [
+        lambda t: t.geo_point.set_srid(4326),
+        lambda t: point_0.set_srid(4326),
+        lambda t: point_1.set_srid(4326),
+        lambda t: point_2,
+        lambda t: point_3.set_srid(4326),
+        lambda t: point_4,
+    ]
+)
+@pytest.mark.only_on_backends(all_db_geo_supported)
+@pytest.mark.xfail_unsupported
+def test_cast_geography(backend, geo, expr_fn):
+    """Testing for geo spatial transform operation."""
+    p = expr_fn(geo).cast('geography')
+    expr = geo[geo.id, p.distance(p).name('tmp')]
+    result = expr.execute()['tmp'][[0]]
+    # distance from a point to a same point should be 0
+    assert np.all(result == 0)
+
+
+@pytest.mark.parametrize(
+    'expr_fn',
+    [
+        lambda t: t.geo_point.set_srid(4326),
+        lambda t: point_0.set_srid(4326),
+        lambda t: point_1.set_srid(4326),
+        lambda t: point_2,
+        lambda t: point_3.set_srid(4326),
+        lambda t: point_4,
+    ]
+)
+@pytest.mark.only_on_backends(all_db_geo_supported)
+@pytest.mark.xfail_unsupported
+def test_cast_geometry(backend, geo, expr_fn):
+    """Testing for geo spatial transform operation."""
+    p = expr_fn(geo).cast('geometry')
+    expr = geo[geo.id, p.distance(p).name('tmp')]
+    result = expr.execute()['tmp'][[0]]
+    # distance from a point to a same point should be 0
+    assert np.all(result == 0)

ibis/tests/backends.py

@@ -1,4 +1,5 @@
 import os
+import sys
 from pathlib import Path
 
 import numpy as np
@@ -234,6 +235,10 @@ def connect(self, data_directory):
             host=host, user=user, password=password, database=database
         )
 
+    def geo(self):
+        if sys.version_info >= (3, 6):
+            return self.db.geo
+
 
 class MapD(Backend):
     check_dtype = False