Add geography serde tests and dataframe api tests for python binding

Author: Kontinuation

python/sedona/core/geom/geography.py

@@ -15,28 +15,11 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-import pickle
-
-from sedona.utils.decorators import require
+from shapely.geometry.base import BaseGeometry
 
 
 class Geography:
+    geometry: BaseGeometry
 
-    def __init__(self, geometry):
-        self._geom = geometry
-        self.userData = None
-
-    def getUserData(self):
-        return self.userData
-
-    @classmethod
-    def from_jvm_instance(cls, java_obj):
-        return Geography(java_obj.geometry)
-
-    @classmethod
-    def serialize_for_java(cls, geogs):
-        return pickle.dumps(geogs)
-
-    @require(["Geography"])
-    def create_jvm_instance(self, jvm):
-        return jvm.Geography(self._geom)
+    def __init__(self, geometry: BaseGeometry):
+        self.geometry = geometry

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

@@ -68,7 +68,7 @@ def sqlType(cls):
         return BinaryType()
 
     def serialize(self, obj):
-        return geometry_serde.serialize(obj._geom)
+        return geometry_serde.serialize(obj.geometry)
 
     def deserialize(self, datum):
         geom, offset = geometry_serde.deserialize(datum)

python/sedona/utils/geometry_adapter.py

@@ -15,28 +15,23 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from typing import Union
-
 from shapely.geometry.base import BaseGeometry
 
 from sedona.core.geom.envelope import Envelope
-from sedona.core.geom.geography import Geography
 from sedona.core.jvm.translate import JvmGeometryAdapter
 from sedona.utils.spatial_rdd_parser import GeometryFactory
 
 
 class GeometryAdapter:
 
     @classmethod
-    def create_jvm_geometry_from_base_geometry(
-        cls, jvm, geom: Union[BaseGeometry, Geography]
-    ):
+    def create_jvm_geometry_from_base_geometry(cls, jvm, geom: BaseGeometry):
         """
         :param jvm:
         :param geom:
         :return:
         """
-        if isinstance(geom, (Envelope, Geography)):
+        if isinstance(geom, Envelope):
             jvm_geom = geom.create_jvm_instance(jvm)
         else:
             decoded_geom = GeometryFactory.to_bytes(geom)

python/tests/sql/test_dataframe_api.py

@@ -26,6 +26,7 @@
 from shapely.geometry.base import BaseGeometry
 from tests.test_base import TestBase
 
+from sedona.core.geom.geography import Geography
 from sedona.sql import st_aggregates as sta
 from sedona.sql import st_constructors as stc
 from sedona.sql import st_functions as stf
@@ -85,6 +86,8 @@
     (stc.ST_GeomFromWKT, ("wkt",), "linestring_wkt", "", "LINESTRING (1 2, 3 4)"),
     (stc.ST_GeomFromWKT, ("wkt", 4326), "linestring_wkt", "", "LINESTRING (1 2, 3 4)"),
     (stc.ST_GeomFromEWKT, ("ewkt",), "linestring_ewkt", "", "LINESTRING (1 2, 3 4)"),
+    (stc.ST_GeogFromWKT, ("wkt",), "linestring_wkt", "", "LINESTRING (1 2, 3 4)"),
+    (stc.ST_GeogFromWKT, ("wkt", 4326), "linestring_wkt", "", "LINESTRING (1 2, 3 4)"),
     (stc.ST_LineFromText, ("wkt",), "linestring_wkt", "", "LINESTRING (1 2, 3 4)"),
     (
         stc.ST_LineFromWKB,
@@ -1230,6 +1233,7 @@
     (stc.ST_LinestringFromWKB, (None,)),
     (stc.ST_GeomFromEWKB, (None,)),
     (stc.ST_GeomFromWKT, (None,)),
+    (stc.ST_GeogFromWKT, (None,)),
     (stc.ST_GeometryFromText, (None,)),
     (stc.ST_LineFromText, (None,)),
     (stc.ST_LineStringFromText, (None, "")),
@@ -1711,6 +1715,9 @@ def test_dataframe_function(
         if isinstance(actual_result, BaseGeometry):
             self.assert_geometry_almost_equal(expected_result, actual_result)
             return
+        elif isinstance(actual_result, Geography):
+            self.assert_geometry_almost_equal(expected_result, actual_result.geometry)
+            return
         elif isinstance(actual_result, bytearray):
             actual_result = actual_result.hex()
         elif isinstance(actual_result, Row):

python/tests/sql/test_geography.py

@@ -0,0 +1,46 @@
+#  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.
+
+import pytest
+from pyspark.sql.functions import expr
+from pyspark.sql.types import StructType
+from shapely.wkt import loads as wkt_loads
+from sedona.core.geom.geography import Geography
+from sedona.sql.types import GeographyType
+from tests.test_base import TestBase
+
+
+class TestGeography(TestBase):
+
+    def test_deserialize_geography(self):
+        """Test serialization and deserialization of Geography objects"""
+        geog_df = self.spark.range(0, 10).withColumn(
+            "geog", expr("ST_GeogFromWKT(CONCAT('POINT (', id, ' ', id + 1, ')'))")
+        )
+        rows = geog_df.collect()
+        assert len(rows) == 10
+        for row in rows:
+            id = row["id"]
+            geog = row["geog"]
+            assert geog.geometry.wkt == f"POINT ({id} {id + 1})"
+
+    def test_serialize_geography(self):
+        wkt = "MULTIPOLYGON (((10 10, 20 20, 20 10, 10 10)), ((-10 -10, -20 -20, -20 -10, -10 -10)))"
+        geog = Geography(wkt_loads(wkt))
+        schema = StructType().add("geog", GeographyType())
+        returned_geog = self.spark.createDataFrame([(geog,)], schema).take(1)[0][0]
+        assert geog.geometry.equals(returned_geog.geometry)

spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala

@@ -2692,7 +2692,9 @@ class functionTestScala
     assert(functionDf.first().get(0) == null)
     functionDf = sparkSession.sql("select ST_AsBinary(null)")
     assert(functionDf.first().get(0) == null)
-    functionDf = sparkSession.sql("select ST_AsEWKB(null)")
+    functionDf = sparkSession.sql("select ST_AsEWKB(ST_GeomFromWKT(null))")
+    assert(functionDf.first().get(0) == null)
+    functionDf = sparkSession.sql("select ST_AsEWKB(ST_GeogFromWKT(null))")
     assert(functionDf.first().get(0) == null)
     functionDf = sparkSession.sql("select ST_SRID(null)")
     assert(functionDf.first().get(0) == null)
@@ -2764,7 +2766,9 @@ class functionTestScala
     assert(functionDf.first().get(0) == null)
     functionDf = sparkSession.sql("select ST_Reverse(null)")
     assert(functionDf.first().get(0) == null)
-    functionDf = sparkSession.sql("select ST_AsEWKT(null)")
+    functionDf = sparkSession.sql("select ST_AsEWKT(ST_GeomFromWKT(null))")
+    assert(functionDf.first().get(0) == null)
+    functionDf = sparkSession.sql("select ST_AsEWKT(ST_GeogFromWKT(null))")
     assert(functionDf.first().get(0) == null)
     functionDf = sparkSession.sql("select ST_Force_2D(null)")
     assert(functionDf.first().get(0) == null)