[SEDONA-635] Allow ST_AsGeoJSON to return feature and featureCollecti…

…on format GeoJSON (#1530) * feat: add type parameter * update docs * update docs - add examples for features and featureCollections * fix typo * fix snowflake spotless error * fix python dataframe api * fix python dataframe api test
apache · Jul 31, 2024 · 9ebdd5d · 9ebdd5d
1 parent a1bcdef
commit 9ebdd5d
Show file tree

Hide file tree

Showing 16 changed files with 322 additions and 19 deletions.
diff --git a/common/src/main/java/org/apache/sedona/common/Functions.java b/common/src/main/java/org/apache/sedona/common/Functions.java
@@ -64,6 +64,8 @@
 import org.locationtech.jts.simplify.VWSimplifier;
 import org.locationtech.jts.triangulate.DelaunayTriangulationBuilder;
 import org.locationtech.jts.triangulate.polygon.ConstrainedDelaunayTriangulator;
+import org.wololo.geojson.Feature;
+import org.wololo.geojson.FeatureCollection;
 import org.wololo.jts2geojson.GeoJSONWriter;
 
 public class Functions {
@@ -664,11 +666,39 @@ public static byte[] asWKB(Geometry geometry) {
   }
 
   public static String asGeoJson(Geometry geometry) {
+    return asGeoJson(geometry, "simple");
+  }
+
+  public static String asGeoJson(Geometry geometry, String type) {
     if (geometry == null) {
       return null;
     }
+
     GeoJSONWriter writer = new GeoJSONWriter();
-    return writer.write(geometry).toString();
+    org.wololo.geojson.Geometry geoJson = writer.write(geometry);
+
+    switch (type.toLowerCase()) {
+      case "simple":
+        return geoJson.toString();
+
+      case "feature":
+        Map<String, Object> properties = new HashMap<>();
+        Feature feature = new Feature(geoJson, properties);
+        return feature.toString();
+
+      case "featurecollection":
+        List<Feature> features = new ArrayList<>();
+        features.add(new Feature(geoJson, new HashMap<>()));
+        FeatureCollection featureCollection =
+            new FeatureCollection(features.toArray(new Feature[0]));
+        return featureCollection.toString();
+
+      default:
+        throw new IllegalArgumentException(
+            "Unknown type: "
+                + type
+                + ". Valid types are: 'simple', 'feature', 'featurecollection'.");
+    }
   }
 
   public static int nPoints(Geometry geometry) {

diff --git a/docs/api/flink/Function.md b/docs/api/flink/Function.md
@@ -372,13 +372,23 @@ POINT ZM(1 1 1 1)
 
 ## ST_AsGeoJSON
 
-Introduction: Return the [GeoJSON](https://geojson.org/) string representation of a geometry
+Introduction: Return the [GeoJSON](https://geojson.org/) string representation of a geometry.
 
-Format: `ST_AsGeoJSON (A: Geometry)`
+The type parameter (Since: `v1.6.1`) takes the following options -
+
+- "Simple" (default): Returns a simple GeoJSON geometry.
+- "Feature": Wraps the geometry in a GeoJSON Feature.
+- "FeatureCollection": Wraps the Feature in a GeoJSON FeatureCollection.
+
+Format:
+
+`ST_AsGeoJSON (A: Geometry)`
+
+`ST_AsGeoJSON (A: Geometry, type: String)`
 
 Since: `v1.3.0`
 
-Example:
+SQL Example (Simple GeoJSON):
 
 ```sql
 SELECT ST_AsGeoJSON(ST_GeomFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'))
@@ -399,6 +409,50 @@ Output:
 }
 ```
 
+SQL Example (Feature GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"Feature",
+  "geometry": {
+      "type":"Polygon",
+      "coordinates":[
+        [[1.0,1.0],
+          [8.0,1.0],
+          [8.0,8.0],
+          [1.0,8.0],
+          [1.0,1.0]]
+      ]
+  }
+}
+```
+
+SQL Example (FeatureCollection GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"FeatureCollection",
+  "features": [{
+    "type":"Feature",
+    "geometry": {
+      "type":"Polygon",
+      "coordinates":[
+        [[1.0,1.0],
+          [8.0,1.0],
+          [8.0,8.0],
+          [1.0,8.0],
+          [1.0,1.0]]
+      ]
+    }
+  }
+  ]
+}
+```
+
 ## ST_AsGML
 
 Introduction: Return the [GML](https://www.ogc.org/standards/gml) string representation of a geometry

diff --git a/docs/api/snowflake/vector-data/Function.md b/docs/api/snowflake/vector-data/Function.md
@@ -270,13 +270,81 @@ FROM polygondf
 
 Introduction: Return the [GeoJSON](https://geojson.org/) string representation of a geometry
 
-Format: `ST_AsGeoJSON (A:geometry)`
+The type parameter takes the following options -
 
-SQL example:
+- "Simple" (default): Returns a simple GeoJSON geometry.
+- "Feature": Wraps the geometry in a GeoJSON Feature.
+- "FeatureCollection": Wraps the Feature in a GeoJSON FeatureCollection.
 
-```SQL
-SELECT ST_AsGeoJSON(polygondf.countyshape)
-FROM polygondf
+Format:
+
+`ST_AsGeoJSON (A:geometry)`
+
+`ST_AsGeoJSON (A:geometry, type: String)`
+
+SQL Example (Simple GeoJSON):
+
+```sql
+SELECT ST_AsGeoJSON(ST_GeomFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'))
+```
+
+Output:
+
+```json
+{
+  "type":"Polygon",
+  "coordinates":[
+    [[1.0,1.0],
+      [8.0,1.0],
+      [8.0,8.0],
+      [1.0,8.0],
+      [1.0,1.0]]
+  ]
+}
+```
+
+SQL Example (Feature GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"Feature",
+  "geometry": {
+      "type":"Polygon",
+      "coordinates":[
+        [[1.0,1.0],
+          [8.0,1.0],
+          [8.0,8.0],
+          [1.0,8.0],
+          [1.0,1.0]]
+      ]
+  }
+}
+```
+
+SQL Example (FeatureCollection GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"FeatureCollection",
+  "features": [{
+      "type":"Feature",
+      "geometry": {
+          "type":"Polygon",
+          "coordinates":[
+            [[1.0,1.0],
+              [8.0,1.0],
+              [8.0,8.0],
+              [1.0,8.0],
+              [1.0,1.0]]
+          ]
+      }
+    }
+  ]
+}
 ```
 
 ## ST_AsGML

diff --git a/docs/api/sql/Function.md b/docs/api/sql/Function.md
@@ -371,11 +371,21 @@ POINT ZM(1 1 1 1)
 
 Introduction: Return the [GeoJSON](https://geojson.org/) string representation of a geometry
 
-Format: `ST_AsGeoJSON (A: Geometry)`
+The type parameter (Since: `v1.6.1`) takes the following options -
+
+- "Simple" (default): Returns a simple GeoJSON geometry.
+- "Feature": Wraps the geometry in a GeoJSON Feature.
+- "FeatureCollection": Wraps the Feature in a GeoJSON FeatureCollection.
+
+Format:
+
+`ST_AsGeoJSON (A: Geometry)`
+
+`ST_AsGeoJSON (A: Geometry, type: String)`
 
 Since: `v1.0.0`
 
-SQL Example
+SQL Example (Simple GeoJSON):
 
 ```sql
 SELECT ST_AsGeoJSON(ST_GeomFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'))
@@ -396,6 +406,50 @@ Output:
 }
 ```
 
+SQL Example (Feature GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"Feature",
+  "geometry": {
+      "type":"Polygon",
+      "coordinates":[
+        [[1.0,1.0],
+          [8.0,1.0],
+          [8.0,8.0],
+          [1.0,8.0],
+          [1.0,1.0]]
+      ]
+  }
+}
+```
+
+SQL Example (FeatureCollection GeoJSON):
+
+Output:
+
+```json
+{
+  "type":"FeatureCollection",
+  "features": [{
+      "type":"Feature",
+      "geometry": {
+          "type":"Polygon",
+          "coordinates":[
+            [[1.0,1.0],
+              [8.0,1.0],
+              [8.0,8.0],
+              [1.0,8.0],
+              [1.0,1.0]]
+          ]
+      }
+    }
+  ]
+}
+```
+
 ## ST_AsGML
 
 Introduction: Return the [GML](https://www.ogc.org/standards/gml) string representation of a geometry

diff --git a/flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java b/flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java
@@ -744,6 +744,14 @@ public String eval(
       Geometry geom = (Geometry) o;
       return org.apache.sedona.common.Functions.asGeoJson(geom);
     }
+
+    @DataTypeHint("String")
+    public String eval(
+        @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o,
+        String type) {
+      Geometry geom = (Geometry) o;
+      return org.apache.sedona.common.Functions.asGeoJson(geom, type);
+    }
   }
 
   public static class ST_AsGML extends ScalarFunction {

diff --git a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -1004,6 +1004,27 @@ public void testGeoJSON() {
     assertEquals(
         "{\"type\":\"Polygon\",\"coordinates\":[[[-0.5,-0.5],[-0.5,0.5],[0.5,0.5],[0.5,-0.5],[-0.5,-0.5]]]}",
         result);
+
+    polygonTable = createPolygonTable(testDataSize);
+    polygonTable =
+        polygonTable.select(
+            call(Functions.ST_AsGeoJSON.class.getSimpleName(), $(polygonColNames[0]), "Feature"));
+    result = (String) first(polygonTable).getField(0);
+    assertEquals(
+        "{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[-0.5,-0.5],[-0.5,0.5],[0.5,0.5],[0.5,-0.5],[-0.5,-0.5]]]},\"properties\":{}}",
+        result);
+
+    polygonTable = createPolygonTable(testDataSize);
+    polygonTable =
+        polygonTable.select(
+            call(
+                Functions.ST_AsGeoJSON.class.getSimpleName(),
+                $(polygonColNames[0]),
+                "FeatureCollection"));
+    result = (String) first(polygonTable).getField(0);
+    assertEquals(
+        "{\"type\":\"FeatureCollection\",\"features\":[{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[-0.5,-0.5],[-0.5,0.5],[0.5,0.5],[0.5,-0.5],[-0.5,-0.5]]]},\"properties\":{}}]}",
+        result);
   }
 
   @Test

diff --git a/python/sedona/sql/st_functions.py b/python/sedona/sql/st_functions.py
@@ -160,15 +160,16 @@ def ST_AsEWKT(geometry: ColumnOrName) -> Column:
 
 
 @validate_argument_types
-def ST_AsGeoJSON(geometry: ColumnOrName) -> Column:
+def ST_AsGeoJSON(geometry: ColumnOrName, type: Optional[Union[ColumnOrName, str]] = None) -> Column:
     """Generate the GeoJSON style representation of a geometry column.
 
     :param geometry: Geometry column to generate GeoJSON for.
     :type geometry: ColumnOrName
     :return: GeoJSON representation of geometry as a string column.
     :rtype: Column
     """
-    return _call_st_function("ST_AsGeoJSON", geometry)
+    args = (geometry) if type is None else (geometry, type)
+    return _call_st_function("ST_AsGeoJSON", args)
 
 
 @validate_argument_types

diff --git a/python/tests/sql/test_dataframe_api.py b/python/tests/sql/test_dataframe_api.py
@@ -92,6 +92,8 @@
     (stf.ST_AsHEXEWKB, ("point",), "point_geom", "", "01010000000000000000000000000000000000F03F"),
     (stf.ST_AsEWKT, (lambda: f.expr("ST_SetSRID(point, 4326)"),), "point_geom", "", "SRID=4326;POINT (0 1)"),
     (stf.ST_AsGeoJSON, ("point",), "point_geom", "", "{\"type\":\"Point\",\"coordinates\":[0.0,1.0]}"),
+    (stf.ST_AsGeoJSON, ("point", lambda: f.lit("feature")), "point_geom", "", "{\"type\":\"Feature\",\"geometry\":{\"type\":\"Point\",\"coordinates\":[0.0,1.0]},\"properties\":{}}"),
+    (stf.ST_AsGeoJSON, ("point", lambda: f.lit("featurecollection")), "point_geom", "", "{\"type\":\"FeatureCollection\",\"features\":[{\"type\":\"Feature\",\"geometry\":{\"type\":\"Point\",\"coordinates\":[0.0,1.0]},\"properties\":{}}]}"),
     (stf.ST_AsGML, ("point",), "point_geom", "", "<gml:Point>\n  <gml:coordinates>\n    0.0,1.0 \n  </gml:coordinates>\n</gml:Point>\n"),
     (stf.ST_AsKML, ("point",), "point_geom", "", "<Point>\n  <coordinates>0.0,1.0</coordinates>\n</Point>\n"),
     (stf.ST_AsText, ("point",), "point_geom", "", "POINT (0 1)"),

diff --git a/snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctions.java b/snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctions.java
@@ -154,6 +154,16 @@ public void test_ST_AsGeoJSON() {
     verifySqlSingleRes(
         "select sedona.ST_AsGeoJSON(sedona.ST_GeomFromText('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'))",
         "{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]}");
+
+    registerUDF("ST_AsGeoJSON", byte[].class, String.class);
+    verifySqlSingleRes(
+        "select sedona.ST_AsGeoJSON(sedona.ST_GeomFromText('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'), 'feature')",
+        "{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]},\"properties\":{}}");
+
+    registerUDF("ST_AsGeoJSON", byte[].class, String.class);
+    verifySqlSingleRes(
+        "select sedona.ST_AsGeoJSON(sedona.ST_GeomFromText('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'), 'featurecollection')",
+        "{\"type\":\"FeatureCollection\",\"features\":[{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]},\"properties\":{}}]}");
   }
 
   @Test

diff --git a/snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctionsV2.java b/snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctionsV2.java
@@ -149,6 +149,16 @@ public void test_ST_AsGeoJSON() {
     verifySqlSingleRes(
         "select sedona.ST_AsGeoJSON(ST_GeometryFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'))",
         "{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]}");
+
+    registerUDFV2("ST_AsGeoJSON", String.class, String.class);
+    verifySqlSingleRes(
+        "select sedona.ST_AsGeoJSON(ST_GeometryFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'), 'feature')",
+        "{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]},\"properties\":{}}");
+
+    registerUDFV2("ST_AsGeoJSON", String.class, String.class);
+    verifySqlSingleRes(
+        "select sedona.ST_AsGeoJSON(ST_GeometryFromWKT('POLYGON((1 1, 8 1, 8 8, 1 8, 1 1))'), 'featurecollection')",
+        "{\"type\":\"FeatureCollection\",\"features\":[{\"type\":\"Feature\",\"geometry\":{\"type\":\"Polygon\",\"coordinates\":[[[1.0,1.0],[8.0,1.0],[8.0,8.0],[1.0,8.0],[1.0,1.0]]]},\"properties\":{}}]}");
   }
 
   @Test