[SEDONA-237] Add ST_Dimension (#867)

Co-authored-by: Jia Yu <jiayu@apache.org>

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

@@ -555,6 +555,14 @@ public static Geometry split(Geometry input, Geometry blade) {
         return new GeometrySplitter(GEOMETRY_FACTORY).split(input, blade);
     }
 
+    public static Integer dimension(Geometry geometry) {
+        Integer dimension = geometry.getDimension();
+        // unknown dimension such as an empty GEOMETRYCOLLECTION
+        if (dimension < 0) {
+            dimension = 0;
+        }
+        return dimension;
+    }
 
     /**
      * get the coordinates of a geometry and transform to Google s2 cell id

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

@@ -267,6 +267,82 @@ public void splitHeterogeneousGeometryCollection() {
         assertNull(actualResult);
     }
 
+    @Test
+    public void dimensionGeometry2D() {
+        Point point = GEOMETRY_FACTORY.createPoint(new Coordinate(1, 2));
+        Integer actualResult = Functions.dimension(point);
+        Integer expectedResult = 0;
+        assertEquals(actualResult, expectedResult);
+
+        LineString lineString = GEOMETRY_FACTORY.createLineString(coordArray(0.0, 0.0, 1.5, 1.5));
+        actualResult = Functions.dimension(lineString);
+        expectedResult = 1;
+        assertEquals(actualResult, expectedResult);
+
+        Polygon polygon = GEOMETRY_FACTORY.createPolygon(coordArray(0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0));
+        actualResult = Functions.dimension(polygon);
+        expectedResult = 2;
+        assertEquals(actualResult, expectedResult);
+
+        MultiPoint multiPoint = GEOMETRY_FACTORY.createMultiPointFromCoords(coordArray(0.0, 0.0, 1.0, 1.0));
+        actualResult = Functions.dimension(multiPoint);
+        expectedResult = 0;
+        assertEquals(actualResult, expectedResult);
+
+        LineString lineString2 = GEOMETRY_FACTORY.createLineString(coordArray(1.0, 1.0, 2.0, 2.0));
+        MultiLineString multiLineString = GEOMETRY_FACTORY
+                .createMultiLineString(new LineString[] { lineString, lineString2 });
+        actualResult = Functions.dimension(multiLineString);
+        expectedResult = 1;
+        assertEquals(actualResult, expectedResult);
+
+        Polygon polygon2 = GEOMETRY_FACTORY.createPolygon(coordArray(0.0, 0.0, 2.0, 2.0, 1.0, 0.0, 0.0, 0.0));
+        MultiPolygon multiPolygon = GEOMETRY_FACTORY.createMultiPolygon(new Polygon[] { polygon, polygon2 });
+        actualResult = Functions.dimension(multiPolygon);
+        expectedResult = 2;
+        assertEquals(actualResult, expectedResult);
+    }
+
+    @Test
+    public void dimensionGeometry3D() {
+        Point point3D = GEOMETRY_FACTORY.createPoint(new Coordinate(1, 1, 1));
+        Integer actualResult = Functions.dimension(point3D);
+        Integer expectedResult = 0;
+        assertEquals(actualResult, expectedResult);
+
+        LineString lineString3D = GEOMETRY_FACTORY.createLineString(coordArray(1, 0, 1, 1, 1, 2));
+        actualResult = Functions.dimension(lineString3D);
+        expectedResult = 1;
+        assertEquals(actualResult, expectedResult);
+
+        Polygon polygon3D = GEOMETRY_FACTORY.createPolygon(coordArray3d(1, 1, 1, 2, 2, 2, 3, 3, 3, 1, 1, 1));
+        actualResult = Functions.dimension(polygon3D);
+        expectedResult = 2;
+        assertEquals(actualResult, expectedResult);
+    }
+
+    @Test
+    public void dimensionGeometryCollection() {
+        Geometry[] geometry = new Geometry[] {
+                GEOMETRY_FACTORY.createLineString(coordArray(0.0, 0.0, 1.5, 1.5)),
+                GEOMETRY_FACTORY.createPolygon(coordArray(0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0))
+        };
+        GeometryCollection geometryCollection = GEOMETRY_FACTORY.createGeometryCollection(geometry);
+
+        Integer actualResult = Functions.dimension(geometryCollection);
+        Integer expectedResult = 2;
+        assertEquals(actualResult, expectedResult);
+    }
+
+    @Test
+    public void dimensionGeometryEmpty() {
+        GeometryCollection emptyGeometryCollection = GEOMETRY_FACTORY.createGeometryCollection();
+
+        Integer actualResult = Functions.dimension(emptyGeometryCollection);
+        Integer expectedResult = 0;
+        assertEquals(actualResult, expectedResult);
+    }
+
     private static boolean intersects(Set<?> s1, Set<?> s2) {
         Set<?> copy = new HashSet<>(s1);
         copy.retainAll(s2);

docs/api/flink/Function.md

@@ -289,6 +289,27 @@ Input: `Polygon ((0 0, 1 2, 2 2, 3 2, 5 0, 4 0, 3 1, 2 1, 1 0, 0 0))`
 
 Output: `POLYGON ((1 2, 2 2, 3 2, 5 0, 4 0, 1 0, 0 0, 1 2))`
 
+## ST_Dimension
+
+Introduction: Return the topological dimension of this Geometry object, which must be less than or equal to the coordinate dimension. OGC SPEC s2.1.1.1 - returns 0 for POINT, 1 for LINESTRING, 2 for POLYGON, and the largest dimension of the components of a GEOMETRYCOLLECTION. If the dimension is unknown (e.g. for an empty GEOMETRYCOLLECTION) 0 is returned.
+
+Format: `ST_Dimension (A:geometry), ST_Dimension (C:geometrycollection), `
+
+Since: `v1.5.0`
+
+Example:
+```sql
+SELECT ST_Dimension('GEOMETRYCOLLECTION(LINESTRING(1 1,0 0),POINT(0 0))');
+```
+
+Result:
+
+```
+ST_Dimension
+-----------
+1
+```
+
 ## ST_Distance
 
 Introduction: Return the Euclidean distance between A and B

docs/api/sql/Function.md

@@ -427,6 +427,28 @@ Result:
 POLYGON ((0 -3, -3 -3, -3 3, 0 3, 0 -3))
 ```
 
+## ST_Dimension
+
+Introduction: Return the topological dimension of this Geometry object, which must be less than or equal to the coordinate dimension. OGC SPEC s2.1.1.1 - returns 0 for POINT, 1 for LINESTRING, 2 for POLYGON, and the largest dimension of the components of a GEOMETRYCOLLECTION. If the dimension is unknown (e.g. for an empty GEOMETRYCOLLECTION) 0 is returned.
+
+Format: `ST_Dimension (A:geometry), ST_Dimension (C:geometrycollection), `
+
+Since: `v1.5.0`
+
+Example:
+```sql
+SELECT ST_Dimension('GEOMETRYCOLLECTION(LINESTRING(1 1,0 0),POINT(0 0))');
+```
+
+Result:
+
+```
+ST_Dimension
+-----------
+1
+```
+
+
 ## ST_Distance
 
 Introduction: Return the Euclidean distance between A and B

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

@@ -43,6 +43,7 @@ public static UserDefinedFunction[] getFuncs() {
                 new Functions.ST_Buffer(),
                 new Functions.ST_ConcaveHull(),
                 new Functions.ST_Envelope(),
+                new Functions.ST_Dimension(),
                 new Functions.ST_Distance(),
                 new Functions.ST_DistanceSphere(),
                 new Functions.ST_DistanceSpheroid(),

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

@@ -17,6 +17,7 @@
 import org.apache.flink.table.annotation.DataTypeHint;
 import org.apache.flink.table.functions.ScalarFunction;
 import org.locationtech.jts.geom.Geometry;
+import org.locationtech.jts.geom.GeometryCollection;
 import org.opengis.referencing.FactoryException;
 import org.opengis.referencing.operation.TransformException;
 
@@ -88,6 +89,15 @@ public Geometry eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.j
         }
     }
 
+    public static class ST_Dimension extends ScalarFunction {
+        @DataTypeHint("Integer")
+        public Integer eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o) {
+            Geometry geom = (Geometry) o;
+            return org.apache.sedona.common.Functions.dimension(geom);
+        }
+
+    }
+
     public static class ST_Distance extends ScalarFunction {
         @DataTypeHint("Double")
         public Double eval(@DataTypeHint(value = "RAW", bridgedTo = org.locationtech.jts.geom.Geometry.class) Object o1,

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

@@ -154,6 +154,16 @@ public void testTransformWKT() throws FactoryException {
     }
 
 
+    @Test
+    public void testDimension(){
+        Table pointTable = tableEnv.sqlQuery(
+                        "SELECT ST_Dimension(ST_GeomFromWKT('GEOMETRYCOLLECTION EMPTY'))");
+        assertEquals(0, first(pointTable).getField(0));
+
+        pointTable = tableEnv.sqlQuery(
+                "SELECT ST_Dimension(ST_GeomFromWKT('GEOMETRYCOLLECTION(MULTIPOLYGON(((0 0, 0 1, 1 1, 1 0, 0 0)), ((2 2, 2 3, 3 3, 3 2, 2 2))), MULTIPOINT(6 6, 7 7, 8 8))'))");
+        assertEquals(2, first(pointTable).getField(0));
+    }
     @Test
     public void testDistance() {
         Table pointTable = createPointTable(testDataSize);

python/sedona/sql/st_functions.py

@@ -45,6 +45,7 @@
     "ST_ConcaveHull",
     "ST_ConvexHull",
     "ST_Difference",
+    "ST_Dimension",
     "ST_Distance",
     "ST_DistanceSphere",
     "ST_DistanceSpheroid",
@@ -385,6 +386,16 @@ def ST_ConvexHull(geometry: ColumnOrName) -> Column:
     """
     return _call_st_function("ST_ConvexHull", geometry)
 
+@validate_argument_types
+def ST_Dimension(geometry: ColumnOrName):
+    """Calculate the inherent dimension of a geometry column.
+
+    :param geometry: Geometry column to calculate the dimension for.
+    :type geometry: ColumnOrName
+    :return: Dimension of geometry as an integer column.
+    :rtype: Column
+    """
+    return _call_st_function("ST_Dimension", geometry)
 
 @validate_argument_types
 def ST_Difference(a: ColumnOrName, b: ColumnOrName) -> Column:

python/tests/sql/test_dataframe_api.py

@@ -75,6 +75,7 @@
     (stf.ST_ConcaveHull, ("geom", 1.0, True), "triangle_geom", "", "POLYGON ((1 1, 1 0, 0 0, 1 1))"),
     (stf.ST_ConvexHull, ("geom",), "triangle_geom", "", "POLYGON ((0 0, 1 1, 1 0, 0 0))"),
     (stf.ST_Difference, ("a", "b"), "overlapping_polys", "", "POLYGON ((1 0, 0 0, 0 1, 1 1, 1 0))"),
+    (stf.ST_Dimension, ("geom",), "geometry_geom_collection", "", 1),
     (stf.ST_Distance, ("a", "b"), "two_points", "", 3.0),
     (stf.ST_DistanceSpheroid, ("point", "point"), "point_geom", "", 0.0),
     (stf.ST_DistanceSphere, ("point", "point"), "point_geom", "", 0.0),
@@ -392,6 +393,8 @@ def base_df(self, request):
             return TestDataFrameAPI.spark.sql("SELECT ST_GeomFromWKT('LINESTRING (0 0, 2 1)') AS line, ST_GeomFromWKT('POLYGON ((1 0, 2 0, 2 2, 1 2, 1 0))') AS poly")
         elif request.param == "square_geom":
             return TestDataFrameAPI.spark.sql("SELECT ST_GeomFromWKT('POLYGON ((1 0, 1 1, 2 1, 2 0, 1 0))') AS geom")
+        elif request.param == "geometry_geom_collection":
+            return TestDataFrameAPI.spark.sql("SELECT ST_GeomFromWKT('GEOMETRYCOLLECTION(POINT(1 1), LINESTRING(0 0, 1 1, 2 2))') AS geom")
         elif request.param == "point_and_line":
             return TestDataFrameAPI.spark.sql("SELECT ST_GeomFromWKT('POINT (0.0 1.0)') AS point, ST_GeomFromWKT('LINESTRING (0 0, 1 0, 2 0, 3 0, 4 0, 5 0)') AS line")
         raise ValueError(f"Invalid base_df name passed: {request.param}")

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

@@ -54,6 +54,7 @@ object Catalog {
     function[ST_Within](),
     function[ST_Covers](),
     function[ST_CoveredBy](),
+    function[ST_Dimension](),
     function[ST_Disjoint](),
     function[ST_Distance](),
     function[ST_3DDistance](),

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

@@ -1010,6 +1010,13 @@ case class ST_Translate(inputExpressions: Seq[Expression])
   }
 }
 
+case class ST_Dimension(inputExpressions: Seq[Expression])
+  extends InferredUnaryExpression(Functions.dimension) with FoldableExpression {
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {
+    copy(inputExpressions = newChildren)
+  }
+}
+
 case class ST_BoundingDiagonal(inputExpressions: Seq[Expression])
   extends InferredUnaryExpression(Functions.boundingDiagonal) with FoldableExpression {
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = {

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

@@ -90,6 +90,9 @@ object st_functions extends DataFrameAPI {
   def ST_Difference(a: Column, b: Column): Column = wrapExpression[ST_Difference](a, b)
   def ST_Difference(a: String, b: String): Column = wrapExpression[ST_Difference](a, b)
 
+  def ST_Dimension(geometry: Column): Column = wrapExpression[ST_Dimension](geometry)
+  def ST_Dimension(geometry: String): Column = wrapExpression[ST_Dimension](geometry)
+
   def ST_Distance(a: Column, b: Column): Column = wrapExpression[ST_Distance](a, b)
   def ST_Distance(a: String, b: String): Column = wrapExpression[ST_Distance](a, b)
 

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

@@ -231,6 +231,14 @@ class dataFrameAPITestScala extends TestBaseScala {
       assert(actualResult == expectedResult)
     }
 
+    it("Passed ST_Dimension") {
+      val polygonDf = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON ((0 0, 1 0, 1 1, 0 0))') AS geom")
+      val df = polygonDf.select(ST_Dimension("geom"))
+      val actualResult = df.take(1)(0).get(0).asInstanceOf[Int]
+      val expectedResult = 2
+      assert(actualResult == expectedResult)
+    }
+
     it("Passed ST_Distance") {
       val pointDf = sparkSession.sql("SELECT ST_Point(0.0, 0.0) AS a, ST_Point(1.0, 0.0) as b")
       val df = pointDf.select(ST_Distance("a", "b"))

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

@@ -137,6 +137,37 @@ class functionTestScala extends TestBaseScala with Matchers with GeometrySample
       assert(functionDf.count() > 0);
     }
 
+    it("Passed ST_Dimension with Geometry") {
+      val geomTestCases = Map(
+        ("'POINT (51.3168 -0.56)'") -> "0",
+        ("'LineString (0 0, 0 90)'") -> "1",
+        ("'POLYGON((0 0,0 5,5 0,0 0))'") -> "2",
+        ("'MULTILINESTRING((0 0, 0 5, 5 0, 0 0))'") -> "1"
+      )
+      for ((geom, expectedResult) <- geomTestCases) {
+        val df = sparkSession.sql(s"SELECT ST_Dimension(ST_GeomFromWKT($geom)), " +
+          s"$expectedResult")
+        val actual = df.take(1)(0).get(0).asInstanceOf[Integer]
+        val expected = df.take(1)(0).get(1).asInstanceOf[Integer]
+        assertEquals(expected, actual)
+      }
+    }
+
+    it("Passed DT_Dimension with GeometryCollection") {
+      val geomTestCases = Map(
+        ("'GEOMETRYCOLLECTION EMPTY'") -> "0",
+        ("'GEOMETRYCOLLECTION(LINESTRING(1 1,0 0),POINT(0 0))'") -> "1",
+        ("'GEOMETRYCOLLECTION(MULTIPOLYGON(((0 0, 0 1, 1 1, 1 0, 0 0)), ((2 2, 2 3, 3 3, 3 2, 2 2))), MULTIPOINT(6 6, 7 7, 8 8))'") -> "2"
+      )
+      for ((geom, expectedResult) <- geomTestCases) {
+        val df = sparkSession.sql(s"SELECT ST_Dimension(ST_GeomFromWKT($geom)), " +
+          s"$expectedResult")
+        val actual = df.take(1)(0).get(0).asInstanceOf[Integer]
+        val expected = df.take(1)(0).get(1).asInstanceOf[Integer]
+        assertEquals(expected, actual)
+      }
+    }
+
     it("Passed ST_Distance") {
       var polygonWktDf = sparkSession.read.format("csv").option("delimiter", "\t").option("header", "false").load(mixedWktGeometryInputLocation)
       polygonWktDf.createOrReplaceTempView("polygontable")