[TASK-78] Add ST_ForceCollection (apache#164)

* feat: Add ST_ForceCollection

* fix: snowflake tests

* fix: snowflake tests

* chore: remove print statement

* chore: add GeometryCollection test

* docs: add detailed explanation

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

@@ -1553,6 +1553,18 @@ public static Geometry force3D(Geometry geometry) {
         return GeomUtils.get3DGeom(geometry, 0.0);
     }
 
+    public static Geometry forceCollection(Geometry geom) {
+        return new GeometryFactory().createGeometryCollection(convertGeometryToArray(geom));
+    }
+
+    private static Geometry[] convertGeometryToArray(Geometry geom) {
+        Geometry[] array = new Geometry[geom.getNumGeometries()];
+        for (int i = 0; i < array.length; i++) {
+            array[i] = geom.getGeometryN(i);
+        }
+        return array;
+    }
+
     public static Integer nRings(Geometry geometry) throws Exception {
         String geometryType = geometry.getGeometryType();
         if (!(geometry instanceof Polygon || geometry instanceof MultiPolygon)) {

common/src/test/java/org/apache/sedona/common/FunctionsTest.java

@@ -1305,6 +1305,34 @@ public void force3DObject3D() {
         assertEquals(expectedDims, Functions.nDims(forcedLine3D));
     }
 
+    @Test
+    public void forceCollection() throws ParseException {
+        Geometry geom = Constructors.geomFromWKT("MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)", 0);
+        int actual = Functions.numGeometries(Functions.forceCollection(geom));
+        int expected = 6;
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("MULTIPOLYGON(((0 0 0,0 1 0,1 1 0,1 0 0,0 0 0)),((0 0 0,1 0 0,1 0 1,0 0 1,0 0 0)),((1 1 0,1 1 1,1 0 1,1 0 0,1 1 0)),((0 1 0,0 1 1,1 1 1,1 1 0,0 1 0)),((0 0 1,1 0 1,1 1 1,0 1 1,0 0 1)))", 0);
+        actual = Functions.numGeometries(Functions.forceCollection(geom));
+        expected = 5;
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("MULTILINESTRING ((10 10, 20 20, 30 30), (15 15, 25 25, 35 35))", 0);
+        actual = Functions.numGeometries(Functions.forceCollection(geom));
+        expected = 2;
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))", 0);
+        actual = Functions.numGeometries(Functions.forceCollection(geom));
+        expected = 1;
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("GEOMETRYCOLLECTION(POLYGON((0 0 1,0 5 1,5 0 1,0 0 1),(1 1 1,3 1 1,1 3 1,1 1 1)))",0);
+        String actualWKT = Functions.asWKT(Functions.forceCollection(geom));
+        String expectedWKT = "GEOMETRYCOLLECTION Z(POLYGON Z((0 0 1, 0 5 1, 5 0 1, 0 0 1), (1 1 1, 3 1 1, 1 3 1, 1 1 1)))";
+        assertEquals(expectedWKT, actualWKT);
+    }
+
     @Test
     public void force3DObject3DDefaultValue() {
         int expectedDims = 3;

docs/api/flink/Function.md

@@ -1253,6 +1253,30 @@ Output:
 LINESTRING EMPTY
 ```
 
+## ST_ForceCollection
+
+Introduction: This function converts the input geometry into a GeometryCollection, regardless of the original geometry type. If the input is a multipart geometry, such as a MultiPolygon or MultiLineString, it will be decomposed into a GeometryCollection containing each individual Polygon or LineString element from the original multipart geometry.
+
+Format: `ST_ForceCollection(geom: Geometry)`
+
+Since: `vTBD`
+
+SQL Example
+
+```sql
+SELECT ST_ForceCollection(
+            ST_GeomFromWKT(
+                "MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)"
+    )
+)
+```
+
+Output:
+
+```
+GEOMETRYCOLLECTION (POINT (30 10), POINT (40 40), POINT (20 20), POINT (10 30), POINT (10 10), POINT (20 50))
+```
+
 ## ST_ForcePolygonCCW
 
 Introduction: For (Multi)Polygon geometries, this function sets the exterior ring orientation to counter-clockwise and interior rings to clockwise orientation. Non-polygonal geometries are returned unchanged.

docs/api/snowflake/vector-data/Function.md

@@ -972,6 +972,28 @@ Input: `LINESTRING EMPTY`
 
 Output: `LINESTRING EMPTY`
 
+## ST_ForceCollection
+
+Introduction: This function converts the input geometry into a GeometryCollection, regardless of the original geometry type. If the input is a multipart geometry, such as a MultiPolygon or MultiLineString, it will be decomposed into a GeometryCollection containing each individual Polygon or LineString element from the original multipart geometry.
+
+Format: `ST_ForceCollection(geom: Geometry)`
+
+SQL Example
+
+```sql
+SELECT ST_ForceCollection(
+            ST_GeomFromWKT(
+                "MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)"
+    )
+)
+```
+
+Output:
+
+```
+GEOMETRYCOLLECTION (POINT (30 10), POINT (40 40), POINT (20 20), POINT (10 30), POINT (10 10), POINT (20 50))
+```
+
 ## ST_ForcePolygonCCW
 
 Introduction: For (Multi)Polygon geometries, this function sets the exterior ring orientation to counter-clockwise and interior rings to clockwise orientation. Non-polygonal geometries are returned unchanged.

docs/api/sql/Function.md

@@ -1260,6 +1260,30 @@ Output:
 LINESTRING EMPTY
 ```
 
+## ST_ForceCollection
+
+Introduction: This function converts the input geometry into a GeometryCollection, regardless of the original geometry type. If the input is a multipart geometry, such as a MultiPolygon or MultiLineString, it will be decomposed into a GeometryCollection containing each individual Polygon or LineString element from the original multipart geometry.
+
+Format: `ST_ForceCollection(geom: Geometry)`
+
+Since: `vTBD`
+
+SQL Example
+
+```sql
+SELECT ST_ForceCollection(
+            ST_GeomFromWKT(
+                "MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)"
+    )
+)
+```
+
+Output:
+
+```
+GEOMETRYCOLLECTION (POINT (30 10), POINT (40 40), POINT (20 20), POINT (10 30), POINT (10 10), POINT (20 50))
+```
+
 ## ST_ForcePolygonCCW
 
 Introduction: For (Multi)Polygon geometries, this function sets the exterior ring orientation to counter-clockwise and interior rings to clockwise orientation. Non-polygonal geometries are returned unchanged.

flink/src/main/java/org/apache/sedona/flink/Catalog.java

@@ -158,6 +158,7 @@ public static UserDefinedFunction[] getFuncs() {
                 new Functions.ST_GeometricMedian(),
                 new Functions.ST_NumPoints(),
                 new Functions.ST_Force3D(),
+                new Functions.ST_ForceCollection(),
                 new Functions.ST_ForcePolygonCW(),
                 new Functions.ST_ForceRHR(),
                 new Functions.ST_NRings(),

flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java

@@ -1102,6 +1102,14 @@ public Geometry eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.j
         }
     }
 
+    public static class ST_ForceCollection 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 o) {
+            Geometry geometry = (Geometry) o;
+            return org.apache.sedona.common.Functions.forceCollection(geometry);
+        }
+    }
+
     public static class ST_ForcePolygonCW 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 o) {

flink/src/test/java/org/apache/sedona/flink/FunctionTest.java

@@ -1314,6 +1314,37 @@ public void testForce3DDefaultValue() {
         assertEquals(expectedDims, actual);
     }
 
+    @Test
+    public void testForceCollection() {
+        int actual = (int) first(
+                tableEnv.sqlQuery("SELECT ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)') AS geom").select(call(Functions.ST_ForceCollection.class.getSimpleName(), $("geom"))).as("geom")
+                        .select(call(Functions.ST_NumGeometries.class.getSimpleName(), $("geom")))
+        ).getField(0);
+        int expected = 6;
+        assertEquals(expected, actual);
+
+        actual = (int) first(
+                tableEnv.sqlQuery("SELECT ST_GeomFromWKT('MULTIPOLYGON(((0 0 0,0 1 0,1 1 0,1 0 0,0 0 0)),((0 0 0,1 0 0,1 0 1,0 0 1,0 0 0)),((1 1 0,1 1 1,1 0 1,1 0 0,1 1 0)),((0 1 0,0 1 1,1 1 1,1 1 0,0 1 0)),((0 0 1,1 0 1,1 1 1,0 1 1,0 0 1)))') AS geom").select(call(Functions.ST_ForceCollection.class.getSimpleName(), $("geom"))).as("geom")
+                        .select(call(Functions.ST_NumGeometries.class.getSimpleName(), $("geom")))
+        ).getField(0);
+        expected = 5;
+        assertEquals(expected, actual);
+
+        actual = (int) first(
+                tableEnv.sqlQuery("SELECT ST_GeomFromWKT('MULTILINESTRING ((10 10, 20 20, 30 30), (15 15, 25 25, 35 35))') AS geom").select(call(Functions.ST_ForceCollection.class.getSimpleName(), $("geom"))).as("geom")
+                        .select(call(Functions.ST_NumGeometries.class.getSimpleName(), $("geom")))
+        ).getField(0);
+        expected = 2;
+        assertEquals(expected, actual);
+
+        actual = (int) first(
+                tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))') AS geom").select(call(Functions.ST_ForceCollection.class.getSimpleName(), $("geom"))).as("geom")
+                        .select(call(Functions.ST_NumGeometries.class.getSimpleName(), $("geom")))
+        ).getField(0);
+        expected = 1;
+        assertEquals(expected, actual);
+    }
+
     @Test
     public void testTriangulatePolygon() {
         Table polyTable = tableEnv.sqlQuery("SELECT ST_TriangulatePolygon(ST_GeomFromWKT('POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0), (5 5, 5 8, 8 8, 8 5, 5 5))')) as poly");

python/sedona/sql/st_functions.py

@@ -1690,6 +1690,16 @@ def ST_Force3D(geometry: ColumnOrName, zValue: Optional[Union[ColumnOrName, floa
     args = (geometry, zValue)
     return _call_st_function("ST_Force3D", args)
 
+@validate_argument_types
+def ST_ForceCollection(geometry: ColumnOrName) -> Column:
+    """
+    Converts a geometry to a geometry collection
+
+    :param geometry: Geometry column to change orientation
+    :return: a Geometry Collection
+    """
+    return _call_st_function("ST_ForceCollection", geometry)
+
 @validate_argument_types
 def ST_ForcePolygonCW(geometry: ColumnOrName) -> Column:
     """

python/tests/sql/test_dataframe_api.py

@@ -111,6 +111,7 @@
     (stf.ST_FlipCoordinates, ("point",), "point_geom", "", "POINT (1 0)"),
     (stf.ST_Force_2D, ("point",), "point_geom", "", "POINT (0 1)"),
     (stf.ST_Force3D, ("point", 1.0), "point_geom", "", "POINT Z (0 1 1)"),
+    (stf.ST_ForceCollection, ("multipoint",), "multipoint_geom", "ST_NumGeometries(geom)", 4),
     (stf.ST_ForcePolygonCW, ("geom",), "geom_with_hole", "", "POLYGON ((0 0, 3 3, 3 0, 0 0), (1 1, 2 1, 2 2, 1 1))"),
     (stf.ST_ForcePolygonCCW, ("geom",), "geom_with_hole", "", "POLYGON ((0 0, 3 0, 3 3, 0 0), (1 1, 2 2, 2 1, 1 1))"),
     (stf.ST_ForceRHR, ("geom",), "geom_with_hole", "", "POLYGON ((0 0, 3 3, 3 0, 0 0), (1 1, 2 1, 2 2, 1 1))"),
@@ -292,6 +293,7 @@
     (stf.ST_ExteriorRing, (None,)),
     (stf.ST_FlipCoordinates, (None,)),
     (stf.ST_Force_2D, (None,)),
+    (stf.ST_ForceCollection, (None,)),
     (stf.ST_ForcePolygonCW, (None,)),
     (stf.ST_ForcePolygonCCW, (None,)),
     (stf.ST_ForceRHR, (None,)),

python/tests/sql/test_function.py

@@ -1468,6 +1468,14 @@ def test_force3D(self):
         actual = actualDf.selectExpr("ST_NDims(geom)").take(1)[0][0]
         assert expected == actual
 
+    def test_st_force_collection(self):
+        basedf = self.spark.sql("SELECT ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)') AS mpoint, ST_GeomFromWKT('POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))') AS poly")
+        actual = basedf.selectExpr("ST_NumGeometries(ST_ForceCollection(mpoint))").take(1)[0][0]
+        assert actual == 6
+
+        actual = basedf.selectExpr("ST_NumGeometries(ST_ForceCollection(poly))").take(1)[0][0]
+        assert actual == 1
+
     def test_forcePolygonCW(self):
         actualDf = self.spark.sql("SELECT ST_ForcePolygonCW(ST_GeomFromWKT('POLYGON ((20 35, 10 30, 10 10, 30 5, 45 20, 20 35),(30 20, 20 15, 20 25, 30 20))')) AS polyCW")
         actual = actualDf.selectExpr("ST_AsText(polyCW)").take(1)[0][0]

snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctions.java

@@ -1072,6 +1072,16 @@ public void test_ST_Force3D() {
         );
     }
 
+    @Test
+    public void test_ST_ForceCollection() {
+        registerUDF("ST_ForceCollection", byte[].class);
+        registerUDF("ST_NumGeometries", byte[].class);
+        verifySqlSingleRes(
+                "SELECT sedona.ST_NumGeometries(sedona.ST_ForceCollection(sedona.ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)')))",
+                6
+        );
+    }
+
     @Test
     public void test_ST_ForcePolygonCW() {
         registerUDF("ST_ForcePolygonCW", byte[].class);

snowflake-tester/src/test/java/org/apache/sedona/snowflake/snowsql/TestFunctionsV2.java

@@ -1031,6 +1031,16 @@ public void test_ST_Force3D() {
         );
     }
 
+    @Test
+    public void test_ST_ForceCollection() {
+        registerUDFV2("ST_ForceCollection", String.class);
+        registerUDFV2("ST_NumGeometries", String.class);
+        verifySqlSingleRes(
+                "SELECT ST_NumGeometries(sedona.ST_ForceCollection(ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)')))",
+                6
+        );
+    }
+
     @Test
     public void test_ST_ForcePolygonCW() {
         registerUDFV2("ST_ForcePolygonCW", String.class);

snowflake/src/main/java/org/apache/sedona/snowflake/snowsql/UDFs.java

@@ -1435,6 +1435,15 @@ public static byte[] ST_Force3D(byte[] geom) {
         );
     }
 
+    @UDFAnnotations.ParamMeta(argNames = {"geom"})
+    public static byte[] ST_ForceCollection(byte[] geom) {
+        return GeometrySerde.serialize(
+                Functions.forceCollection(
+                        GeometrySerde.deserialize(geom)
+                )
+        );
+    }
+
     @UDFAnnotations.ParamMeta(argNames = {"geom"})
     public static byte[] ST_ForcePolygonCW(byte[] geom) {
         return GeometrySerde.serialize(

snowflake/src/main/java/org/apache/sedona/snowflake/snowsql/UDFsV2.java

@@ -1217,6 +1217,15 @@ public static String ST_Force3D(String geom) {
         );
     }
 
+    @UDFAnnotations.ParamMeta(argNames = {"geom"}, argTypes = {"Geometry"}, returnTypes = "Geometry")
+    public static String ST_ForceCollection(String geom) {
+        return GeometrySerde.serGeoJson(
+                Functions.forceCollection(
+                        GeometrySerde.deserGeoJson(geom)
+                )
+        );
+    }
+
     @UDFAnnotations.ParamMeta(argNames = {"geom"}, argTypes = {"Geometry"}, returnTypes = "Geometry")
     public static String ST_ForcePolygonCW(String geom) {
         return GeometrySerde.serGeoJson(

spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala

@@ -193,6 +193,7 @@ object Catalog {
     function[ST_LengthSpheroid](),
     function[ST_NumPoints](),
     function[ST_Force3D](0.0),
+    function[ST_ForceCollection](),
     function[ST_NRings](),
     function[ST_Translate](0.0),
     function[ST_TriangulatePolygon](),

spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala

@@ -1200,6 +1200,15 @@ case class ST_Force3D(inputExpressions: Seq[Expression])
   }
 }
 
+case class ST_ForceCollection(inputExpressions: Seq[Expression])
+  extends InferredExpression(Functions.forceCollection _) {
+
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
+    copy(inputExpressions = newChildren)
+  }
+}
+
+
 case class ST_ForcePolygonCW(inputExpressions: Seq[Expression])
   extends InferredExpression(Functions.forcePolygonCW _) {
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {

spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala

@@ -427,6 +427,10 @@ object st_functions extends DataFrameAPI {
 
   def ST_Force3D(geometry: String, zValue: Double): Column = wrapExpression[ST_Force3D](geometry, zValue)
 
+  def ST_ForceCollection(geometry: Column): Column = wrapExpression[ST_ForceCollection](geometry)
+
+  def ST_ForceCollection(geometry: String): Column = wrapExpression[ST_ForceCollection](geometry)
+
   def ST_ForcePolygonCW(geometry: Column): Column = wrapExpression[ST_ForcePolygonCW](geometry)
   def ST_ForcePolygonCW(geometry: String): Column = wrapExpression[ST_ForcePolygonCW](geometry)
 

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

@@ -1348,6 +1348,15 @@ class dataFrameAPITestScala extends TestBaseScala {
       assertEquals(expectedGeomDefaultValue, wktWriter.write(actualGeomDefaultValue))
     }
 
+    it("Passed ST_ForceCollection") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)') AS mpoint, ST_GeomFromWKT('POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))') AS poly")
+      var actual = baseDf.select(ST_NumGeometries(ST_ForceCollection("mpoint"))).first().get(0)
+      assert(actual == 6)
+
+      actual = baseDf.select(ST_NumGeometries(ST_ForceCollection("poly"))).first().get(0)
+      assert(actual == 1)
+    }
+
     it("Passed ST_TriangulatePolygon") {
       val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0), (5 5, 5 8, 8 8, 8 5, 5 5))') as poly")
       val actual = baseDf.select(ST_AsText(ST_TriangulatePolygon("poly"))).first().getString(0)

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

@@ -2274,6 +2274,14 @@ class functionTestScala extends TestBaseScala with Matchers with GeometrySample
     }
   }
 
+  it("Passed ST_ForceCollection") {
+    var actual = sparkSession.sql("SELECT ST_NumGeometries(ST_ForceCollection(ST_GeomFromWKT('MULTIPOINT (30 10, 40 40, 20 20, 10 30, 10 10, 20 50)')))").first().get(0)
+    assert(actual == 6)
+
+    actual = sparkSession.sql("SELECT ST_NumGeometries(ST_ForceCollection(ST_GeomFromWKT('POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))')))").first().get(0)
+    assert(actual == 1)
+  }
+
   it("should pass ST_NRings") {
     val geomTestCases = Map(
       ("'POLYGON ((1 0, 1 1, 2 1, 2 0, 1 0))'") -> 1,