[SEDONA-303] Port all Sedona Spark function to Sedona Flink -- ST_Col…

…lect (#902)
apache · Jul 14, 2023 · 58f7692 · 58f7692
1 parent 9108472
commit 58f7692
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 0 deletions.
diff --git a/docs/api/flink/Function.md b/docs/api/flink/Function.md
@@ -471,6 +471,58 @@ Input: `g1: 'POLYGON ((190 150, 20 10, 160 70, 190 150))', g2: ST_Buffer('POINT(
 
 Output: `POINT(131.59149149528952 101.89887534906197)`
 
+## ST_Collect
+
+Introduction: Returns MultiGeometry object based on geometry column/s or array with geometries
+
+Format
+
+`ST_Collect(*geom: geometry)`
+
+`ST_Collect(geom: array<geometry>)`
+
+Since: `v1.5.0`
+
+Example:
+
+```sql
+SELECT ST_Collect(
+    ST_GeomFromText('POINT(21.427834 52.042576573)'),
+    ST_GeomFromText('POINT(45.342524 56.342354355)')
+) AS geom
+```
+
+Result:
+
+```
++---------------------------------------------------------------+
+|geom                                                           |
++---------------------------------------------------------------+
+|MULTIPOINT ((21.427834 52.042576573), (45.342524 56.342354355))|
++---------------------------------------------------------------+
+```
+
+Example:
+
+```sql
+SELECT ST_Collect(
+    Array(
+        ST_GeomFromText('POINT(21.427834 52.042576573)'),
+        ST_GeomFromText('POINT(45.342524 56.342354355)')
+    )
+) AS geom
+```
+
+Result:
+
+```
++---------------------------------------------------------------+
+|geom                                                           |
++---------------------------------------------------------------+
+|MULTIPOINT ((21.427834 52.042576573), (45.342524 56.342354355))|
++---------------------------------------------------------------+
+```
+
 ## ST_CollectionExtract
 
 Introduction: Returns a homogeneous multi-geometry from a given geometry collection.

diff --git a/flink/src/main/java/org/apache/sedona/flink/Catalog.java b/flink/src/main/java/org/apache/sedona/flink/Catalog.java
@@ -47,6 +47,7 @@ public static UserDefinedFunction[] getFuncs() {
                 new Functions.ST_Buffer(),
                 new Functions.ST_ClosestPoint(),
                 new Functions.ST_Centroid(),
+                new Functions.ST_Collect(),
                 new Functions.ST_CollectionExtract(),
                 new Functions.ST_ConcaveHull(),
                 new Functions.ST_ConvexHull(),

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
@@ -91,6 +91,23 @@ public Geometry eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.j
             return org.apache.sedona.common.Functions.getCentroid(geom);
         }
     }
+
+    public static class ST_Collect extends ScalarFunction {
+        @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class)
+        public Geometry eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o1, 
+                             @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o2  ) {
+            Geometry geom1 = (Geometry) o1;
+            Geometry geom2 = (Geometry) o2;
+            Geometry[] geoms = new Geometry[]{geom1, geom2};
+            return org.apache.sedona.common.Functions.createMultiGeometry(geoms);
+        }
+
+        @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class)
+        public Geometry eval(@DataTypeHint(inputGroup = InputGroup.ANY) Object o) {
+            Geometry[] geoms = (Geometry[]) o;
+            return org.apache.sedona.common.Functions.createMultiGeometry(geoms);
+        }
+    }
 
     public static class ST_CollectionExtract extends ScalarFunction {
         @DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class)

diff --git a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -100,6 +100,32 @@ public void testCentroid() {
         assertEquals("POINT (1 1)", result.toString());
     }
 
+    @Test
+    public void testCollectWithTwoInputs() {
+        Table pointTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POINT (1 2)') AS g1, ST_GeomFromWKT('POINT (-2 3)') as g2");
+        Table resultTable = pointTable.select(call(Functions.ST_Collect.class.getSimpleName(), $("g1"), $("g2")));
+        Geometry result1 = (Geometry) first(resultTable).getField(0);
+        assertEquals("MULTIPOINT ((1 2), (-2 3))", result1.toString());
+
+        Table collectionTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POINT (1 2)') AS g1, ST_GeomFromWKT('LINESTRING(1 2, 3 4)') as g2");
+        resultTable = collectionTable.select(call(Functions.ST_Collect.class.getSimpleName(), $("g1"), $("g2")));
+        Geometry result2 = (Geometry) first(resultTable).getField(0);
+        assertEquals("GEOMETRYCOLLECTION (POINT (1 2), LINESTRING (1 2, 3 4))", result2.toString());
+    }
+
+    @Test
+    public void testCollectWithArray() {
+        Table lineTable = tableEnv.sqlQuery("SELECT array[ST_GeomFromText('LINESTRING(1 2, 3 4)'), ST_GeomFromText('LINESTRING(3 4, 4 5)')] as lines");
+        Table resultTable = lineTable.select(call(Functions.ST_Collect.class.getSimpleName(), $("lines")));
+        Geometry result1 = (Geometry) first(resultTable).getField(0);
+        assertEquals("MULTILINESTRING ((1 2, 3 4), (3 4, 4 5))", result1.toString());
+
+        Table collectionTable = tableEnv.sqlQuery("SELECT array[ST_GeomFromText('POINT(0 0)'), ST_GeomFromText('LINESTRING(3 4, 4 5)')] as lines");
+        resultTable = collectionTable.select(call(Functions.ST_Collect.class.getSimpleName(), $("lines")));
+        Geometry result2 = (Geometry) first(resultTable).getField(0);
+        assertEquals("GEOMETRYCOLLECTION (POINT (0 0), LINESTRING (3 4, 4 5))", result2.toString());
+    }
+
     @Test
     public void testCollectionExtract() {
         Table collectionTable = tableEnv.sqlQuery("SELECT ST_GeomFromText('GEOMETRYCOLLECTION(POINT(0 0), LINESTRING(1 1, 2 2))') as collection");