[SEDONA-293] Implement ST_IsCollection (apache#910)

Kontinuation · Jul 24, 2023 · 2e24568 · 2e24568
1 parent 8831903
commit 2e24568
Show file tree

Hide file tree

Showing 14 changed files with 160 additions and 0 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
@@ -1079,6 +1079,14 @@ public static double frechetDistance(Geometry g1, Geometry g2) {
         return GeomUtils.getFrechetDistance(g1, g2);
     }
 
+    public static boolean isCollection(Geometry geometry) {
+        String geoType = geometry.getGeometryType();
+        return Geometry.TYPENAME_GEOMETRYCOLLECTION.equalsIgnoreCase(geoType) ||
+                Geometry.TYPENAME_MULTIPOINT.equalsIgnoreCase(geoType) ||
+                Geometry.TYPENAME_MULTIPOLYGON.equalsIgnoreCase(geoType) ||
+                Geometry.TYPENAME_MULTILINESTRING.equalsIgnoreCase(geoType);
+    }
+
     public static Geometry boundingDiagonal(Geometry geometry) {
         if (geometry.isEmpty()) {
             return GEOMETRY_FACTORY.createLineString();

diff --git a/common/src/test/java/org/apache/sedona/common/FunctionsTest.java b/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
@@ -1259,6 +1259,27 @@ public void angleInvalidUnsupportedGeom() {
         assertEquals("ST_Angle supports either only POINT or only LINESTRING geometries.", e.getMessage());
     }
 
+    @Test
+    public void isCollectionWithCollection() {
+        Point[] points = new Point[]{
+                GEOMETRY_FACTORY.createPoint(new Coordinate(5.0, 6.0)),
+                GEOMETRY_FACTORY.createPoint(new Coordinate(7.0, 8.0))
+        };
+        GeometryCollection geometryCollection = GEOMETRY_FACTORY.createGeometryCollection(points);
+
+        boolean actualResult = Functions.isCollection(geometryCollection);
+        boolean expectedResult = true;
+        assertEquals(actualResult, expectedResult);
+    }
+
+    @Test
+    public void isCollectionWithOutCollection() {
+        Point point = GEOMETRY_FACTORY.createPoint(new Coordinate(6.0, 6.0));
+
+        boolean actualResult = Functions.isCollection(point);
+        boolean expectedResult = false;
+        assertEquals(actualResult, expectedResult);
+    }
 
     public void affineEmpty3D() {
         LineString emptyLineString = GEOMETRY_FACTORY.createLineString();

diff --git a/docs/api/flink/Function.md b/docs/api/flink/Function.md
@@ -1184,6 +1184,34 @@ Example:
 SELECT ST_IsClosed(ST_GeomFromText('LINESTRING(0 0, 1 1, 1 0)'))
 ```
 
+## ST_IsCollection
+
+Introduction: Returns `TRUE` if the geometry type of the input is a geometry collection type. 
+Collection types are the following:
+
+- GEOMETRYCOLLECTION
+- MULTI{POINT, POLYGON, LINESTRING}
+
+Format: `ST_IsCollection(geom: geometry)`
+
+Since: `v1.5.0`
+
+Example:
+
+```sql
+SELECT ST_IsCollection(ST_GeomFromText('MULTIPOINT(0 0), (6 6)'))
+```
+
+Output: `true`
+
+Example:
+
+```sql
+SELECT ST_IsCollection(ST_GeomFromText('POINT(5 5)'))
+```
+
+Output: `false`
+
 ## ST_IsEmpty
 
 Introduction: Test if a geometry is empty geometry

diff --git a/docs/api/sql/Function.md b/docs/api/sql/Function.md
@@ -1204,6 +1204,33 @@ SELECT ST_IsClosed(ST_GeomFromText('LINESTRING(0 0, 1 1, 1 0)'))
 
 Output: `false`
 
+## ST_IsCollection
+
+Introduction: Returns `TRUE` if the geometry type of the input is a geometry collection type.
+Collection types are the following:
+
+- GEOMETRYCOLLECTION
+- MULTI{POINT, POLYGON, LINESTRING}
+
+Format: `ST_IsCollection(geom: geometry)`
+
+Since: `v1.5.0`
+
+Example:
+
+```sql
+SELECT ST_IsCollection(ST_GeomFromText('MULTIPOINT(0 0), (6 6)'))
+```
+
+Output: `true`
+
+Example:
+```sql
+SELECT ST_IsCollection(ST_GeomFromText('POINT(5 5)'))
+```
+
+Output: `false`
+
 ## ST_IsEmpty
 
 Introduction: Test if a geometry is empty geometry

diff --git a/flink/src/main/java/org/apache/sedona/flink/Catalog.java b/flink/src/main/java/org/apache/sedona/flink/Catalog.java
@@ -138,6 +138,7 @@ public static UserDefinedFunction[] getFuncs() {
                 new Functions.ST_Angle(),
                 new Functions.ST_Degrees(),
                 new Functions.ST_HausdorffDistance(),
+                new Functions.ST_IsCollection(),
                 new Functions.ST_CoordDim(),
         };
     }

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
@@ -978,6 +978,14 @@ public Integer eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jt
         }
     }
 
+    public static class ST_IsCollection extends ScalarFunction {
+        @DataTypeHint("Boolean")
+        public boolean eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o) {
+            Geometry geom = (Geometry) o;
+            return org.apache.sedona.common.Functions.isCollection(geom);
+        }
+    }
+
     public static class ST_Angle extends ScalarFunction {
 
         @DataTypeHint("Double")

diff --git a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -1095,6 +1095,22 @@ public void testHausdorffDistance() {
         assertEquals(expectedDefault, actualDefault);
     }
 
+    @Test
+    public void testIsCollectionForCollection() {
+        Table collectionTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('GEOMETRYCOLLECTION(POINT(2 3), POINT(4 6), LINESTRING(15 15, 20 20))') AS collection");
+        Table resultTable = collectionTable.select(call(Functions.ST_IsCollection.class.getSimpleName(), $("collection")));
+        boolean result = (boolean) first(resultTable).getField(0);
+        assertTrue(result);
+    }
+
+    @Test
+    public void testIsCollectionForNotCollection() {
+        Table collectionTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POINT(10 10)') AS collection");
+        Table resultTable = collectionTable.select(call(Functions.ST_IsCollection.class.getSimpleName(), $("collection")));
+        boolean result = (boolean) first(resultTable).getField(0);
+        assertFalse(result);
+    }
+
     @Test
     public void testCoordDimFor2D() {
         Table polygonTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POINT(3 7)') AS " + polygonColNames[0]);

diff --git a/python/sedona/sql/st_functions.py b/python/sedona/sql/st_functions.py
@@ -123,6 +123,7 @@
     "ST_Degrees",
     "ST_FrechetDistance",
     "ST_CoordDim",
+    "ST_IsCollection",
     "ST_Affine",
     "ST_BoundingDiagonal"
 ]
@@ -1495,3 +1496,13 @@ def ST_CoordDim(geometry: ColumnOrName) -> Column:
     """
     return _call_st_function("ST_CoordDim", geometry)
 
+@validate_argument_types
+def ST_IsCollection(geometry: ColumnOrName) -> Column:
+    """Check if the geometry is of GeometryCollection type.
+
+    :param geometry: Column for geometry collection
+    :type geometry: ColumnOrName
+    :return: True if geometry is a collection of geometries.
+    :rtype: Column
+    """
+    return _call_st_function("ST_IsCollection", geometry)
diff --git a/python/tests/sql/test_dataframe_api.py b/python/tests/sql/test_dataframe_api.py
@@ -107,6 +107,7 @@
     (stf.ST_HausdorffDistance, ("point", "line",), "point_and_line", "", 5.0990195135927845),
     (stf.ST_InteriorRingN, ("geom", 0), "geom_with_hole", "", "LINESTRING (1 1, 2 2, 2 1, 1 1)"),
     (stf.ST_Intersection, ("a", "b"), "overlapping_polys", "", "POLYGON ((2 0, 1 0, 1 1, 2 1, 2 0))"),
+    (stf.ST_IsCollection, ("geom",), "geom_collection", "", True),
     (stf.ST_IsClosed, ("geom",), "closed_linestring_geom", "", True),
     (stf.ST_IsEmpty, ("geom",), "empty_geom", "", True),
     (stf.ST_IsRing, ("line",), "linestring_geom", "", False),

diff --git a/sql/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala b/sql/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala
@@ -115,6 +115,7 @@ object Catalog {
     function[ST_Dump](),
     function[ST_DumpPoints](),
     function[ST_IsClosed](),
+    function[ST_IsCollection](),
     function[ST_NumInteriorRings](),
     function[ST_AddPoint](-1),
     function[ST_RemovePoint](-1),

diff --git a/sql/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala b/sql/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
@@ -1114,3 +1114,15 @@ case class ST_CoordDim(inputExpressions: Seq[Expression])
   }
 }
 
+/**
+ * Returns True if geometry is a collection of geometries
+ *
+ * @param inputExpressions
+ */
+case class ST_IsCollection(inputExpressions: Seq[Expression])
+  extends InferredExpression(Functions.isCollection _) {
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = {
+    copy(inputExpressions = newChildren)
+  }
+}
+
diff --git a/sql/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala b/sql/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala
@@ -400,5 +400,9 @@ object st_functions extends DataFrameAPI {
 
   def ST_CoordDim(geometry: String): Column = wrapExpression[ST_CoordDim](geometry)
 
+  def ST_IsCollection(geometry: Column): Column = wrapExpression[ST_IsCollection](geometry)
+
+  def ST_IsCollection(geometry: String): Column = wrapExpression[ST_IsCollection](geometry)
+
 }
 
diff --git a/sql/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala b/sql/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala
@@ -1107,6 +1107,22 @@ class dataFrameAPITestScala extends TestBaseScala {
       assert(expected == actual)
     }
 
+    it("Passed ST_IsCollection with a collection") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('GEOMETRYCOLLECTION(POINT(3 2), POINT(6 4))') AS collection")
+      val df = baseDf.select(ST_IsCollection("collection"))
+      val actualResult = df.take(1)(0).getBoolean(0)
+      val expectedResult = true
+      assert(actualResult == expectedResult)
+    }
+
+    it("Passed ST_IsCollection without a collection") {
+      val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('POINT(9 9)') AS collection")
+      val df = baseDf.select(ST_IsCollection("collection"))
+      val actualResult = df.take(1)(0).getBoolean(0)
+      val expectedResult = false
+      assert(actualResult == expectedResult)
+    }
+
     it("Passed ST_Angle - 4 Points") {
       val polyDf = sparkSession.sql("SELECT ST_GeomFromWKT('POINT (10 10)') AS p1, ST_GeomFromWKT('POINT (0 0)') AS p2," +
         " ST_GeomFromWKT('POINT (90 90)') AS p3, ST_GeomFromWKT('POINT (100 80)') AS p4")

diff --git a/sql/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala b/sql/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala
@@ -2176,6 +2176,12 @@ class functionTestScala extends TestBaseScala with Matchers with GeometrySample
     assert(test.take(1)(0).get(0).asInstanceOf[Int] == 4)
   }
 
+  it("Should pass ST_IsCollection") {
+    val test = sparkSession.sql("SELECT ST_IsCollection(ST_GeomFromText('GEOMETRYCOLLECTION(POINT(2 3), POINT(4 6), LINESTRING(15 15, 20 20))'))")
+    val expected = true
+    assert(test.take(1)(0).get(0).asInstanceOf[Boolean] == expected)
+  }
+
   it ("should pass GeometryType") {
     val geomTestCases = Map (
       ("'POINT (51.3168 -0.56)'") -> "'POINT'",