This is an automated email from the ASF dual-hosted git repository.

jiayu pushed a commit to branch SEDONA-584
in repository https://gitbox.apache.org/repos/asf/sedona.git

commit 3ce7c4da376e550ac8c44e04b440a9a0fa89d4a1
Author: Furqaan Khan <[email protected]>
AuthorDate: Thu Apr 11 00:21:57 2024 -0400

    [TASK-46] Add ST_Zmflag (#163)
    
    * fix: typo
    
    * fix: testing snowflake behavior
    
    * fix: testing snowflake behavior
    
    * fix: remove function from snowflake
---
 .../java/org/apache/sedona/common/Functions.java   | 18 ++++++++++
 .../org/apache/sedona/common/FunctionsTest.java    | 16 +++++++++
 docs/api/flink/Function.md                         | 38 ++++++++++++++++++++++
 docs/api/sql/Function.md                           | 38 ++++++++++++++++++++++
 .../main/java/org/apache/sedona/flink/Catalog.java |  1 +
 .../apache/sedona/flink/expressions/Functions.java |  8 +++++
 .../java/org/apache/sedona/flink/FunctionTest.java | 19 +++++++++++
 python/sedona/sql/st_functions.py                  | 12 +++++++
 python/tests/sql/test_dataframe_api.py             |  2 ++
 python/tests/sql/test_function.py                  | 13 ++++++++
 .../scala/org/apache/sedona/sql/UDF/Catalog.scala  |  1 +
 .../sql/sedona_sql/expressions/Functions.scala     |  8 +++++
 .../sql/sedona_sql/expressions/st_functions.scala  |  3 ++
 .../apache/sedona/sql/dataFrameAPITestScala.scala  | 14 ++++++++
 .../org/apache/sedona/sql/functionTestScala.scala  | 14 ++++++++
 15 files changed, 205 insertions(+)

diff --git a/common/src/main/java/org/apache/sedona/common/Functions.java 
b/common/src/main/java/org/apache/sedona/common/Functions.java
index d574dc8fd..6ee9d49ef 100644
--- a/common/src/main/java/org/apache/sedona/common/Functions.java
+++ b/common/src/main/java/org/apache/sedona/common/Functions.java
@@ -762,6 +762,24 @@ public class Functions {
         }
     }
 
+    public static int zmFlag(Geometry geom) {
+        Coordinate coords = geom.getCoordinate();
+        boolean hasZ = !Double.isNaN(coords.getZ());
+        boolean hasM = !Double.isNaN(coords.getM());
+        if (hasM && hasZ) {
+            // geom is 4D
+            return 3;
+        } else if (hasZ) {
+            // geom is 3D-Z
+            return 2;
+        } else if (hasM) {
+            // geom is 3D-M
+            return 1;
+        }
+        // geom is 2D
+        return 0;
+    }
+
     public static Geometry concaveHull(Geometry geometry, double pctConvex, 
boolean allowHoles){
         ConcaveHull concave_hull = new ConcaveHull(geometry);
         concave_hull.setMaximumEdgeLengthRatio(pctConvex);
diff --git a/common/src/test/java/org/apache/sedona/common/FunctionsTest.java 
b/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
index 78c82ba77..beeeb9919 100644
--- a/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
+++ b/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
@@ -2304,6 +2304,22 @@ public class FunctionsTest extends TestBase {
         assertEquals(expected, e2.getMessage());
     }
 
+    @Test
+    public void testZmFlag() throws ParseException {
+        int _2D = 0, _3DM = 1, _3DZ = 2, _4D = 3;
+        Geometry geom = Constructors.geomFromWKT("POINT (1 2)", 0);
+        assertEquals(_2D, Functions.zmFlag(geom));
+
+        geom = Constructors.geomFromWKT("LINESTRING (1 2 3, 4 5 6)", 0);
+        assertEquals(_3DZ, Functions.zmFlag(geom));
+
+        geom = Constructors.geomFromWKT("POLYGON M((1 2 3, 3 4 3, 5 6 3, 3 4 
3, 1 2 3))", 0);
+        assertEquals(_3DM, Functions.zmFlag(geom));
+
+        geom = Constructors.geomFromWKT("MULTIPOLYGON ZM (((30 10 5 1, 40 40 
10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1)), ((15 5 3 1, 20 10 6 2, 10 10 7 3, 15 
5 3 1)))", 0);
+        assertEquals(_4D, Functions.zmFlag(geom));
+    }
+
     @Test
     public void hausdorffDistanceDefaultGeom2D() throws Exception {
         Polygon polygon1 = GEOMETRY_FACTORY.createPolygon(coordArray3d(1, 0, 
1, 1, 1, 2, 2, 1, 5, 2, 0, 1, 1, 0, 1));
diff --git a/docs/api/flink/Function.md b/docs/api/flink/Function.md
index 8139a090f..7efd4db21 100644
--- a/docs/api/flink/Function.md
+++ b/docs/api/flink/Function.md
@@ -3367,3 +3367,41 @@ Output:
 ```
 4.0
 ```
+
+## ST_Zmflag
+
+Introduction: Returns a code indicating the Z and M coordinate dimensions 
present in the input geometry.
+
+Values are: 0 = 2D, 1 = 3D-M, 2 = 3D-Z, 3 = 4D.
+
+Format: `ST_Zmflag(geom: Geometry)`
+
+Since: `vTBD`
+
+SQL Example
+
+```sql
+SELECT ST_Zmflag(
+        ST_GeomFromWKT('LINESTRING Z(1 2 3, 4 5 6)')
+)
+```
+
+Output:
+
+```
+2
+```
+
+SQL Example
+
+```sql
+SELECT ST_Zmflag(
+        ST_GeomFromWKT('POINT ZM(1 2 3 4)')
+)
+```
+
+Output:
+
+```
+3
+```
diff --git a/docs/api/sql/Function.md b/docs/api/sql/Function.md
index 182d2afbf..61060bd15 100644
--- a/docs/api/sql/Function.md
+++ b/docs/api/sql/Function.md
@@ -3477,3 +3477,41 @@ Output:
 ```
 4.0
 ```
+
+## ST_Zmflag
+
+Introduction: Returns a code indicating the Z and M coordinate dimensions 
present in the input geometry.
+
+Values are: 0 = 2D, 1 = 3D-M, 2 = 3D-Z, 3 = 4D.
+
+Format: `ST_Zmflag(geom: Geometry)`
+
+Since: `vTBD`
+
+SQL Example
+
+```sql
+SELECT ST_Zmflag(
+        ST_GeomFromWKT('LINESTRING Z(1 2 3, 4 5 6)')
+)
+```
+
+Output:
+
+```
+2
+```
+
+SQL Example
+
+```sql
+SELECT ST_Zmflag(
+        ST_GeomFromWKT('POINT ZM(1 2 3 4)')
+)
+```
+
+Output:
+
+```
+3
+```
diff --git a/flink/src/main/java/org/apache/sedona/flink/Catalog.java 
b/flink/src/main/java/org/apache/sedona/flink/Catalog.java
index 443fbe932..acbaa04f6 100644
--- a/flink/src/main/java/org/apache/sedona/flink/Catalog.java
+++ b/flink/src/main/java/org/apache/sedona/flink/Catalog.java
@@ -107,6 +107,7 @@ public class Catalog {
                 new Functions.ST_X(),
                 new Functions.ST_Y(),
                 new Functions.ST_Z(),
+                new Functions.ST_Zmflag(),
                 new Functions.ST_YMax(),
                 new Functions.ST_YMin(),
                 new Functions.ST_XMax(),
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
index 0f7593181..99deabe85 100644
--- a/flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java
+++ b/flink/src/main/java/org/apache/sedona/flink/expressions/Functions.java
@@ -602,6 +602,14 @@ public class Functions {
         }
     }
 
+    public static class ST_Zmflag 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.zmFlag(geom);
+        }
+    }
+
     public static class ST_XMax extends ScalarFunction {
         @DataTypeHint("Double")
         public Double eval(@DataTypeHint(value = "RAW", bridgedTo = 
org.locationtech.jts.geom.Geometry.class) Object o) {
diff --git a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java 
b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
index 57dc92ece..ed7408e07 100644
--- a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
+++ b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -709,6 +709,25 @@ public class FunctionTest extends TestBase{
         assertEquals(7.89, first(pointTable).getField(0));
     }
 
+    @Test
+    public void testZmflag() {
+        Table table = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POINT (1 2)') 
AS geom");
+        int actual = (int) 
first(table.select(call(Functions.ST_Zmflag.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(0, actual);
+
+        table = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('LINESTRING (1 2 3, 4 
5 6)') AS geom");
+        actual = (int) 
first(table.select(call(Functions.ST_Zmflag.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(2, actual);
+
+        table = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POLYGON M((1 2 3, 3 
4 3, 5 6 3, 3 4 3, 1 2 3))') AS geom");
+        actual = (int) 
first(table.select(call(Functions.ST_Zmflag.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(1, actual);
+
+        table = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('MULTIPOLYGON ZM 
(((30 10 5 1, 40 40 10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1)), ((15 5 3 1, 20 
10 6 2, 10 10 7 3, 15 5 3 1)))') AS geom");
+        actual = (int) 
first(table.select(call(Functions.ST_Zmflag.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(3, actual);
+    }
+
     @Test
     public void testHasZ() {
         Table polyTable = tableEnv.sqlQuery("SELECT ST_GeomFromWKT('POLYGON ZM 
((30 10 5 1, 40 40 10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1))') as poly");
diff --git a/python/sedona/sql/st_functions.py 
b/python/sedona/sql/st_functions.py
index 02a3001c8..0e7cdbc10 100644
--- a/python/sedona/sql/st_functions.py
+++ b/python/sedona/sql/st_functions.py
@@ -1509,6 +1509,18 @@ def ST_Z(point: ColumnOrName) -> Column:
     """
     return _call_st_function("ST_Z", point)
 
+@validate_argument_types
+def ST_Zmflag(geom: ColumnOrName) -> Column:
+    """Return the code indicating the ZM coordinate dimension of a geometry
+        2D = 0, 3D-M = 1, 3D-Z = 2, 4D = 3
+
+    :param geom: Geometry column
+    :type geom: ColumnOrName
+    :return: Code for coordinate dimension of the geometry as an integer 
column.
+    :rtype: Column
+    """
+    return _call_st_function("ST_Zmflag", geom)
+
 @validate_argument_types
 def ST_ZMax(geometry: ColumnOrName) -> Column:
     """Return the maximum Z coordinate of a geometry.
diff --git a/python/tests/sql/test_dataframe_api.py 
b/python/tests/sql/test_dataframe_api.py
index bd75bf16a..fd38a8a13 100644
--- a/python/tests/sql/test_dataframe_api.py
+++ b/python/tests/sql/test_dataframe_api.py
@@ -202,6 +202,7 @@ test_configurations = [
     (stf.ST_YMax, ("geom",), "triangle_geom", "", 1.0),
     (stf.ST_YMin, ("geom",), "triangle_geom", "", 0.0),
     (stf.ST_Z, ("b",), "two_points", "", 4.0),
+    (stf.ST_Zmflag, ("b",), "two_points", "", 2),
     (stf.ST_IsValidReason, ("geom",), "triangle_geom", "", "Valid Geometry"),
     (stf.ST_IsValidReason, ("geom", 1), "triangle_geom", "", "Valid Geometry"),
 
@@ -386,6 +387,7 @@ wrong_type_configurations = [
     (stf.ST_YMax, (None,)),
     (stf.ST_YMin, (None,)),
     (stf.ST_Z, (None,)),
+    (stf.ST_Zmflag, (None,)),
 
     # predicates
     (stp.ST_Contains, (None, "")),
diff --git a/python/tests/sql/test_function.py 
b/python/tests/sql/test_function.py
index 2bf755825..4e807a40f 100644
--- a/python/tests/sql/test_function.py
+++ b/python/tests/sql/test_function.py
@@ -540,6 +540,19 @@ class TestPredicateJoin(TestBase):
 
         assert (not polygons.count())
 
+    def test_st_zmflag(self):
+        actual = self.spark.sql("SELECT ST_Zmflag(ST_GeomFromWKT('POINT (1 
2)'))").take(1)[0][0]
+        assert actual == 0
+
+        actual = self.spark.sql("SELECT ST_Zmflag(ST_GeomFromWKT('LINESTRING 
(1 2 3, 4 5 6)'))").take(1)[0][0]
+        assert actual == 2
+
+        actual = self.spark.sql("SELECT ST_Zmflag(ST_GeomFromWKT('POLYGON M((1 
2 3, 3 4 3, 5 6 3, 3 4 3, 1 2 3))'))").take(1)[0][0]
+        assert actual == 1
+
+        actual = self.spark.sql("SELECT ST_Zmflag(ST_GeomFromWKT('MULTIPOLYGON 
ZM (((30 10 5 1, 40 40 10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1)), ((15 5 3 1, 
20 10 6 2, 10 10 7 3, 15 5 3 1)))'))").take(1)[0][0]
+        assert actual == 3
+
     def test_st_z_max(self):
         linestring_df = self.spark.sql("SELECT ST_GeomFromWKT('LINESTRING Z (0 
0 1, 0 1 2)') as geom")
         linestring_row = [lnstr_row[0] for lnstr_row in 
linestring_df.selectExpr("ST_ZMax(geom)").collect()]
diff --git 
a/spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala 
b/spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala
index 92c9ac807..33c1bd345 100644
--- a/spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala
+++ b/spark/common/src/main/scala/org/apache/sedona/sql/UDF/Catalog.scala
@@ -112,6 +112,7 @@ object Catalog {
     function[ST_X](),
     function[ST_Y](),
     function[ST_Z](),
+    function[ST_Zmflag](),
     function[ST_StartPoint](),
     function[ST_Snap](),
     function[ST_ClosestPoint](),
diff --git 
a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
 
b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
index 2db5e6166..e02ff47b0 100644
--- 
a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
+++ 
b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Functions.scala
@@ -445,6 +445,14 @@ case class ST_Z(inputExpressions: Seq[Expression])
   }
 }
 
+case class ST_Zmflag(inputExpressions: Seq[Expression])
+  extends InferredExpression(Functions.zmFlag _) {
+
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = 
{
+    copy(inputExpressions = newChildren)
+  }
+}
+
 case class ST_StartPoint(inputExpressions: Seq[Expression])
   extends InferredExpression(Functions.startPoint _) {
 
diff --git 
a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala
 
b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala
index 6dbde20c9..653c787a6 100644
--- 
a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala
+++ 
b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_functions.scala
@@ -377,6 +377,9 @@ object st_functions extends DataFrameAPI {
   def ST_Z(point: Column): Column = wrapExpression[ST_Z](point)
   def ST_Z(point: String): Column = wrapExpression[ST_Z](point)
 
+  def ST_Zmflag(geom: Column): Column = wrapExpression[ST_Zmflag](geom)
+  def ST_Zmflag(geom: String): Column = wrapExpression[ST_Zmflag](geom)
+
   def ST_ZMax(geometry: Column): Column = wrapExpression[ST_ZMax](geometry)
   def ST_ZMax(geometry: String): Column = wrapExpression[ST_ZMax](geometry)
 
diff --git 
a/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala 
b/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala
index 2548b1d3b..067e98e67 100644
--- 
a/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala
+++ 
b/spark/common/src/test/scala/org/apache/sedona/sql/dataFrameAPITestScala.scala
@@ -727,6 +727,20 @@ class dataFrameAPITestScala extends TestBaseScala {
       assert(actualResult == expectedResult)
     }
 
+    it("Should pass ST_Zmflag") {
+      var actual = sparkSession.sql("SELECT ST_GeomFromWKT('POINT (1 2)') AS 
geom").select(ST_Zmflag("geom")).first().get(0)
+      assert(actual == 0)
+
+      actual = sparkSession.sql("SELECT ST_GeomFromWKT('LINESTRING (1 2 3, 4 5 
6)') AS geom").select(ST_Zmflag("geom")).first().get(0)
+      assert(actual == 2)
+
+      actual = sparkSession.sql("SELECT ST_GeomFromWKT('POLYGON M((1 2 3, 3 4 
3, 5 6 3, 3 4 3, 1 2 3))') AS geom").select(ST_Zmflag("geom")).first().get(0)
+      assert(actual == 1)
+
+      actual = sparkSession.sql("SELECT ST_GeomFromWKT('MULTIPOLYGON ZM (((30 
10 5 1, 40 40 10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1)), ((15 5 3 1, 20 10 6 2, 
10 10 7 3, 15 5 3 1)))') AS geom").select(ST_Zmflag("geom")).first().get(0)
+      assert(actual == 3)
+    }
+
     it("Passed ST_StartPoint") {
       val baseDf = sparkSession.sql("SELECT ST_GeomFromWKT('LINESTRING (0 0, 1 
0)') AS geom")
       val df = baseDf.select(ST_StartPoint("geom"))
diff --git 
a/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala 
b/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala
index 8711b3fb2..8c44dc168 100644
--- a/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala
+++ b/spark/common/src/test/scala/org/apache/sedona/sql/functionTestScala.scala
@@ -951,6 +951,20 @@ class functionTestScala extends TestBaseScala with 
Matchers with GeometrySample
 
     }
 
+    it("Passed ST_Zmflag") {
+      var actual = sparkSession.sql("SELECT ST_Zmflag(ST_GeomFromWKT('POINT (1 
2)'))").first().get(0)
+      assert(actual == 0)
+
+      actual = sparkSession.sql("SELECT ST_Zmflag(ST_GeomFromWKT('LINESTRING 
(1 2 3, 4 5 6)'))").first().get(0)
+      assert(actual == 2)
+
+      actual = sparkSession.sql("SELECT ST_Zmflag(ST_GeomFromWKT('POLYGON M((1 
2 3, 3 4 3, 5 6 3, 3 4 3, 1 2 3))'))").first().get(0)
+      assert(actual == 1)
+
+      actual = sparkSession.sql("SELECT ST_Zmflag(ST_GeomFromWKT('MULTIPOLYGON 
ZM (((30 10 5 1, 40 40 10 2, 20 40 15 3, 10 20 20 4, 30 10 5 1)), ((15 5 3 1, 
20 10 6 2, 10 10 7 3, 15 5 3 1)))'))").first().get(0)
+      assert(actual == 3)
+    }
+
     it("Should pass ST_StartPoint function") {
       Given("Polygon Data Frame, Point DataFrame, LineString Data Frame")
 

Reply via email to