[SEDONA-711] Add Geography user-defined type (#1828)

* start basic object

* compiling udt class

* maybe actually register

* add one input and one output function

* maybe one input and one output function

* maybe builds

* some possible Python requirements

* format

* maybe a few more references to geography

* remove word

* Making ST_GeogFromWKT, ST_AsEWKT, ST_AsEWKB work properly

* Add geography serde tests and dataframe api tests for python binding

---------

Co-authored-by: Dewey Dunnington <[email protected]>

Author: Kontinuation

common/src/main/java/org/apache/sedona/common/Constructors.java

@@ -22,6 +22,7 @@
 import javax.xml.parsers.ParserConfigurationException;
 import org.apache.sedona.common.enums.FileDataSplitter;
 import org.apache.sedona.common.enums.GeometryType;
+import org.apache.sedona.common.geometryObjects.Geography;
 import org.apache.sedona.common.utils.FormatUtils;
 import org.apache.sedona.common.utils.GeoHashDecoder;
 import org.locationtech.jts.geom.*;
@@ -44,6 +45,10 @@ public static Geometry geomFromWKT(String wkt, int srid) throws ParseException {
     return new WKTReader(geometryFactory).read(wkt);
   }
 
+  public static Geography geogFromWKT(String wkt, int srid) throws ParseException {
+    return new Geography(geomFromWKT(wkt, srid));
+  }
+
   public static Geometry geomFromEWKT(String ewkt) throws ParseException {
     if (ewkt == null) {
       return null;

common/src/main/java/org/apache/sedona/common/Functions.java

@@ -29,6 +29,7 @@
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.sedona.common.geometryObjects.Circle;
+import org.apache.sedona.common.geometryObjects.Geography;
 import org.apache.sedona.common.sphere.Spheroid;
 import org.apache.sedona.common.subDivide.GeometrySubDivider;
 import org.apache.sedona.common.utils.*;
@@ -776,6 +777,10 @@ public static String asEWKT(Geometry geometry) {
     return GeomUtils.getEWKT(geometry);
   }
 
+  public static String asEWKT(Geography geography) {
+    return asEWKT(geography.getGeometry());
+  }
+
   public static String asWKT(Geometry geometry) {
     return GeomUtils.getWKT(geometry);
   }
@@ -784,6 +789,10 @@ public static byte[] asEWKB(Geometry geometry) {
     return GeomUtils.getEWKB(geometry);
   }
 
+  public static byte[] asEWKB(Geography geography) {
+    return asEWKB(geography.getGeometry());
+  }
+
   public static String asHexEWKB(Geometry geom, String endian) {
     if (endian.equalsIgnoreCase("NDR")) {
       return GeomUtils.getHexEWKB(geom, ByteOrderValues.LITTLE_ENDIAN);

common/src/main/java/org/apache/sedona/common/geometryObjects/Geography.java

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.sedona.common.geometryObjects;
+
+import org.locationtech.jts.geom.Geometry;
+
+public class Geography {
+  private final Geometry geometry;
+
+  public Geography(Geometry geometry) {
+    this.geometry = geometry;
+  }
+
+  public Geometry getGeometry() {
+    return this.geometry;
+  }
+
+  public String toString() {
+    return this.geometry.toText();
+  }
+}

common/src/main/java/org/apache/sedona/common/geometrySerde/GeometrySerde.java

@@ -25,6 +25,7 @@
 import com.esotericsoftware.kryo.io.Output;
 import java.io.Serializable;
 import org.apache.sedona.common.geometryObjects.Circle;
+import org.apache.sedona.common.geometryObjects.Geography;
 import org.locationtech.jts.geom.Envelope;
 import org.locationtech.jts.geom.Geometry;
 import org.locationtech.jts.geom.GeometryCollection;
@@ -36,7 +37,7 @@
  * Provides methods to efficiently serialize and deserialize geometry types.
  *
  * <p>Supports Point, LineString, Polygon, MultiPoint, MultiLineString, MultiPolygon,
- * GeometryCollection, Circle and Envelope types.
+ * GeometryCollection, Circle, Envelope, and Geography types.
  *
  * <p>First byte contains {@link Type#id}. Then go type-specific bytes, followed by user-data
  * attached to the geometry.
@@ -63,6 +64,9 @@ public void write(Kryo kryo, Output out, Object object) {
       out.writeDouble(envelope.getMaxX());
       out.writeDouble(envelope.getMinY());
       out.writeDouble(envelope.getMaxY());
+    } else if (object instanceof Geography) {
+      writeType(out, Type.GEOGRAPHY);
+      writeGeometry(kryo, out, ((Geography) object).getGeometry());
     } else {
       throw new UnsupportedOperationException(
           "Cannot serialize object of type " + object.getClass().getName());
@@ -118,6 +122,10 @@ public Object read(Kryo kryo, Input input, Class aClass) {
             return new Envelope();
           }
         }
+      case GEOGRAPHY:
+        {
+          return new Geography(readGeometry(kryo, input));
+        }
       default:
         throw new UnsupportedOperationException(
             "Cannot deserialize object of type " + geometryType);
@@ -145,7 +153,8 @@ private Geometry readGeometry(Kryo kryo, Input input) {
   private enum Type {
     SHAPE(0),
     CIRCLE(1),
-    ENVELOPE(2);
+    ENVELOPE(2),
+    GEOGRAPHY(3);
 
     private final int id;
 

python/sedona/core/geom/geography.py

@@ -0,0 +1,25 @@
+#  Licensed to the Apache Software Foundation (ASF) under one
+#  or more contributor license agreements.  See the NOTICE file
+#  distributed with this work for additional information
+#  regarding copyright ownership.  The ASF licenses this file
+#  to you under the Apache License, Version 2.0 (the
+#  "License"); you may not use this file except in compliance
+#  with the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from shapely.geometry.base import BaseGeometry
+
+
+class Geography:
+    geometry: BaseGeometry
+
+    def __init__(self, geometry: BaseGeometry):
+        self.geometry = geometry

python/sedona/register/java_libs.py

@@ -26,6 +26,7 @@ class SedonaJvmLib(Enum):
     KNNQuery = "org.apache.sedona.core.spatialOperator.KNNQuery"
     RangeQuery = "org.apache.sedona.core.spatialOperator.RangeQuery"
     Envelope = "org.locationtech.jts.geom.Envelope"
+    Geography = "org.apache.sedona.common.geometryObjects.Geography"
     GeoSerializerData = (
         "org.apache.sedona.python.wrapper.adapters.GeoSparkPythonConverter"
     )

python/sedona/spark/__init__.py

@@ -42,7 +42,7 @@
 from sedona.sql.st_constructors import *
 from sedona.sql.st_functions import *
 from sedona.sql.st_predicates import *
-from sedona.sql.types import GeometryType, RasterType
+from sedona.sql.types import GeometryType, GeographyType, RasterType
 from sedona.utils import KryoSerializer, SedonaKryoRegistrator
 from sedona.utils.adapter import Adapter
 from sedona.utils.geoarrow import dataframe_to_arrow

python/sedona/sql/st_constructors.py

@@ -176,6 +176,22 @@ def ST_GeomFromWKT(
     return _call_constructor_function("ST_GeomFromWKT", args)
 
 
+@validate_argument_types
+def ST_GeogFromWKT(
+    wkt: ColumnOrName, srid: Optional[ColumnOrNameOrNumber] = None
+) -> Column:
+    """Generate a geography column from a Well-Known Text (WKT) string column.
+
+    :param wkt: WKT string column to generate from.
+    :type wkt: ColumnOrName
+    :return: Geography column representing the WKT string.
+    :rtype: Column
+    """
+    args = (wkt) if srid is None else (wkt, srid)
+
+    return _call_constructor_function("ST_GeogFromWKT", args)
+
+
 @validate_argument_types
 def ST_GeomFromEWKT(ewkt: ColumnOrName) -> Column:
     """Generate a geometry column from a OGC Extended Well-Known Text (WKT) string column.

python/sedona/sql/types.py

@@ -33,6 +33,7 @@
     SedonaRaster = None
 
 from ..utils import geometry_serde
+from ..core.geom.geography import Geography
 
 
 class GeometryType(UserDefinedType):
@@ -60,6 +61,31 @@ def scalaUDT(cls):
         return "org.apache.spark.sql.sedona_sql.UDT.GeometryUDT"
 
 
+class GeographyType(UserDefinedType):
+
+    @classmethod
+    def sqlType(cls):
+        return BinaryType()
+
+    def serialize(self, obj):
+        return geometry_serde.serialize(obj.geometry)
+
+    def deserialize(self, datum):
+        geom, offset = geometry_serde.deserialize(datum)
+        return Geography(geom)
+
+    @classmethod
+    def module(cls):
+        return "sedona.sql.types"
+
+    def needConversion(self):
+        return True
+
+    @classmethod
+    def scalaUDT(cls):
+        return "org.apache.spark.sql.sedona_sql.UDT.GeographyUDT"
+
+
 class RasterType(UserDefinedType):
 
     @classmethod

python/sedona/utils/prep.py

@@ -28,6 +28,8 @@
 )
 from shapely.geometry.base import BaseGeometry
 
+from ..core.geom.geography import Geography
+
 
 def assign_all() -> bool:
     geoms = [
@@ -41,6 +43,7 @@ def assign_all() -> bool:
     ]
     assign_udt_shapely_objects(geoms=geoms)
     assign_user_data_to_shapely_objects(geoms=geoms)
+    assign_udt_geography()
     return True
 
 
@@ -55,3 +58,9 @@ def assign_udt_shapely_objects(geoms: List[type(BaseGeometry)]) -> bool:
 def assign_user_data_to_shapely_objects(geoms: List[type(BaseGeometry)]) -> bool:
     for geom in geoms:
         geom.getUserData = lambda geom_instance: geom_instance.userData
+
+
+def assign_udt_geography():
+    from sedona.sql.types import GeographyType
+
+    Geography.__UDT__ = GeographyType()