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

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

commit 89796d57f96fa043900c61fd089a1ae1f59eb2e2
Author: Furqaan Khan <[email protected]>
AuthorDate: Fri Apr 19 01:01:07 2024 -0400

    [TASK-64] Add ST_Force3DM (#166)
    
    * feat: add ST_Force3DM
    
    * fix: lint and remove experimental snowflake porting
    
    * refactor: GeometryForce3DMTransformer to a more concise implementation
    
    * fix: docs typo and add test to verify M ordinate value
---
 .../java/org/apache/sedona/common/Functions.java   | 14 ++++---
 .../common/utils/GeometryForce3DMTransformer.java  | 49 ++++++++++++++++++++++
 .../org/apache/sedona/common/FunctionsTest.java    | 47 +++++++++++++++++++++
 docs/api/flink/Function.md                         | 47 +++++++++++++++++++++
 docs/api/sql/Function.md                           | 47 +++++++++++++++++++++
 .../main/java/org/apache/sedona/flink/Catalog.java |  1 +
 .../apache/sedona/flink/expressions/Functions.java | 16 +++++++
 .../java/org/apache/sedona/flink/FunctionTest.java | 11 +++++
 python/sedona/sql/st_functions.py                  | 11 +++++
 python/tests/sql/test_dataframe_api.py             |  2 +
 python/tests/sql/test_function.py                  |  5 +++
 .../scala/org/apache/sedona/sql/UDF/Catalog.scala  |  1 +
 .../sql/sedona_sql/expressions/Functions.scala     |  8 ++++
 .../sql/sedona_sql/expressions/st_functions.scala  |  5 +++
 .../apache/sedona/sql/dataFrameAPITestScala.scala  | 11 +++++
 .../org/apache/sedona/sql/functionTestScala.scala  | 18 ++++++++
 16 files changed, 288 insertions(+), 5 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
index ac88ce04f..95747364e 100644
--- a/common/src/main/java/org/apache/sedona/common/Functions.java
+++ b/common/src/main/java/org/apache/sedona/common/Functions.java
@@ -21,11 +21,7 @@ import org.apache.commons.lang3.tuple.Pair;
 import org.apache.sedona.common.geometryObjects.Circle;
 import org.apache.sedona.common.sphere.Spheroid;
 import org.apache.sedona.common.subDivide.GeometrySubDivider;
-import org.apache.sedona.common.utils.GeomUtils;
-import org.apache.sedona.common.utils.GeometryGeoHashEncoder;
-import org.apache.sedona.common.utils.GeometrySplitter;
-import org.apache.sedona.common.utils.H3Utils;
-import org.apache.sedona.common.utils.S2Utils;
+import org.apache.sedona.common.utils.*;
 import org.locationtech.jts.algorithm.MinimumBoundingCircle;
 import org.locationtech.jts.algorithm.Orientation;
 import org.locationtech.jts.algorithm.hull.ConcaveHull;
@@ -1545,6 +1541,14 @@ public class Functions {
         return geometry.getNumPoints();
     }
 
+    public static Geometry force3DM(Geometry geom, double mValue) {
+        return GeometryForce3DMTransformer.transform(geom, mValue);
+    }
+
+    public static Geometry force3DM(Geometry geom) {
+        return force3DM(geom, 0.0);
+    }
+
     public static Geometry force3D(Geometry geometry, double zValue) {
         return GeomUtils.get3DGeom(geometry, zValue);
     }
diff --git 
a/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DMTransformer.java
 
b/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DMTransformer.java
new file mode 100644
index 000000000..bd76ee76a
--- /dev/null
+++ 
b/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DMTransformer.java
@@ -0,0 +1,49 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.sedona.common.utils;
+
+import org.apache.sedona.common.Functions;
+import org.locationtech.jts.geom.*;
+import org.locationtech.jts.geom.impl.PackedCoordinateSequenceFactory;
+import org.locationtech.jts.geom.util.GeometryTransformer;
+
+import java.util.Arrays;
+
+public class GeometryForce3DMTransformer extends GeometryTransformer {
+
+    private final double mValue;
+
+    public GeometryForce3DMTransformer(double mValue) {
+        this.mValue = mValue;
+    }
+
+    @Override
+    protected CoordinateSequence transformCoordinates(CoordinateSequence 
coords, Geometry parent) {
+        CoordinateXYM[] newCoords = new CoordinateXYM[coords.size()];
+        for (int i = 0; i < coords.size(); i++) {
+            Coordinate coordinate = coords.getCoordinate(i);
+            newCoords[i] = new CoordinateXYM(coordinate.getX(), 
coordinate.getY(), mValue);
+        }
+
+        return createCoordinateSequence(newCoords);
+    }
+
+    public static Geometry transform(Geometry geometry, double mValue) {
+        if (geometry.getCoordinates().length == 0) return geometry;
+        if (Functions.hasM(geometry)) return geometry;
+
+        GeometryForce3DMTransformer transformer = new 
GeometryForce3DMTransformer(mValue);
+        return transformer.transform(geometry);
+    }
+}
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 017f65e64..08589d09b 100644
--- a/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
+++ b/common/src/test/java/org/apache/sedona/common/FunctionsTest.java
@@ -1333,6 +1333,53 @@ public class FunctionsTest extends TestBase {
         assertEquals(expectedWKT, actualWKT);
     }
 
+    @Test
+    public void testForce3DM() throws ParseException {
+        Geometry geom = Constructors.geomFromWKT("POINT (1 2)",0);
+        String actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        String expected = "POINT M(1 2 5)";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("MULTIPOINT ((1 2), (2 3))",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "MULTIPOINT M((1 2 5), (2 3 5))";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("LINESTRING (1 2, 2 3, 3 4)",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "LINESTRING M(1 2 5, 2 3 5, 3 4 5)";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("MULTILINESTRING ((10 10, 20 20, 30 
30), (15 15, 25 25, 35 35))",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "MULTILINESTRING M((10 10 5, 20 20 5, 30 30 5), (15 15 5, 
25 25 5, 35 35 5))";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("LINEARRING (30 10, 40 40, 20 40, 10 
20, 30 10)",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "LINEARRING M(30 10 5, 40 40 5, 20 40 5, 10 20 5, 30 10 5)";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("POLYGON ((0 0, 10 0, 10 10, 0 10, 0 
0), (4 4, 4 6, 6 6, 6 4, 4 4))",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "POLYGON M((0 0 5, 10 0 5, 10 10 5, 0 10 5, 0 0 5), (4 4 5, 
4 6 5, 6 6 5, 6 4 5, 4 4 5))";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("MULTIPOLYGON (((30 10, 40 40, 20 40, 
10 20, 30 10)), ((15 5, 10 20, 20 30, 15 5)))",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "MULTIPOLYGON M(((30 10 5, 40 40 5, 20 40 5, 10 20 5, 30 10 
5)), ((15 5 5, 10 20 5, 20 30 5, 15 5 5)))";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("GEOMETRYCOLLECTION (POINT (10 10), 
LINESTRING (15 15, 25 25, 35 35), POLYGON ((30 10, 40 40, 20 40, 10 20, 30 
10)))",0);
+        actual = Functions.asWKT(Functions.force3DM(geom, 5));
+        expected = "GEOMETRYCOLLECTION M(POINT M(10 10 5), LINESTRING M(15 15 
5, 25 25 5, 35 35 5), POLYGON M((30 10 5, 40 40 5, 20 40 5, 10 20 5, 30 10 
5)))";
+        assertEquals(expected, actual);
+
+        geom = Constructors.geomFromWKT("POLYGON M((0 0 3, 0 5 3, 5 0 3, 0 0 
3), (1 1 3, 3 1 3, 1 3 3, 1 1 3))", 0);
+        Geometry actualGeom = Functions.force3DM(geom, 10);
+        assertTrue(Predicates.equals(geom, actualGeom));
+    }
+
     @Test
     public void force3DObject3DDefaultValue() {
         int expectedDims = 3;
diff --git a/docs/api/flink/Function.md b/docs/api/flink/Function.md
index 9a51cae82..f00e9de37 100644
--- a/docs/api/flink/Function.md
+++ b/docs/api/flink/Function.md
@@ -1266,6 +1266,53 @@ Output:
 LINESTRING EMPTY
 ```
 
+## ST_Force3DM
+
+Introduction: Forces the geometry into XYM mode. Retains any existing M 
coordinate, but removes the Z coordinate if present. Assigns a default M value 
of 0.0 if `mValue` is not specified.
+
+!!!Note
+Example output is after calling ST_AsText() on returned geometry, which adds M 
for in the WKT.
+
+Format: `ST_Force3DM(geometry: Geometry, mValue: Double = 0.0)`
+
+Since: `vTDB`
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('POLYGON M((0 0 3,0 5 3,5 0 3,0 0 
3),(1 1 3,3 1 3,1 3 3,1 1 3))'), 2.3))
+```
+
+Output:
+
+```
+POLYGON M((0 0 3, 0 5 3, 5 0 3, 0 0 3), (1 1 3, 3 1 3, 1 3 3, 1 1 3))
+```
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('LINESTRING(0 1,1 0,2 0)'), 2.3))
+```
+
+Output:
+
+```
+LINESTRING M(0 1 2.3, 1 0 2.3, 2 0 2.3)
+```
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('LINESTRING Z(0 1 3,1 0 3,2 0 
3)'), 5))
+```
+
+Output:
+
+```
+LINESTRING M(0 1 5, 1 0 5, 2 0 5)
+```
+
 ## ST_Force3DZ
 
 Introduction: Forces the geometry into a 3-dimensional model so that all 
output representations will have X, Y and Z coordinates.
diff --git a/docs/api/sql/Function.md b/docs/api/sql/Function.md
index 9a7a3eafe..b5e3f6224 100644
--- a/docs/api/sql/Function.md
+++ b/docs/api/sql/Function.md
@@ -1271,6 +1271,53 @@ Output:
 LINESTRING EMPTY
 ```
 
+## ST_Force3DM
+
+Introduction: Forces the geometry into XYM mode. Retains any existing M 
coordinate, but removes the Z coordinate if present. Assigns a default M value 
of 0.0 if `mValue` is not specified.
+
+!!!Note
+    Example output is after calling ST_AsText() on returned geometry, which 
adds M for in the WKT.
+
+Format: `ST_Force3DM(geometry: Geometry, mValue: Double = 0.0)`
+
+Since: `vTBD`
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('POLYGON M((0 0 3,0 5 3,5 0 3,0 0 
3),(1 1 3,3 1 3,1 3 3,1 1 3))'), 2.3))
+```
+
+Output:
+
+```
+POLYGON M((0 0 3, 0 5 3, 5 0 3, 0 0 3), (1 1 3, 3 1 3, 1 3 3, 1 1 3))
+```
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('LINESTRING(0 1,1 0,2 0)'), 2.3))
+```
+
+Output:
+
+```
+LINESTRING M(0 1 2.3, 1 0 2.3, 2 0 2.3)
+```
+
+SQL Example
+
+```sql
+SELECT ST_AsText(ST_Force3DM(ST_GeomFromText('LINESTRING Z(0 1 3,1 0 3,2 0 
3)'), 5))
+```
+
+Output:
+
+```
+LINESTRING M(0 1 5, 1 0 5, 2 0 5)
+```
+
 ## ST_Force3DZ
 
 Introduction: Forces the geometry into a 3-dimensional model so that all 
output representations will have X, Y and Z coordinates.
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 df8c86186..a569d4a99 100644
--- a/flink/src/main/java/org/apache/sedona/flink/Catalog.java
+++ b/flink/src/main/java/org/apache/sedona/flink/Catalog.java
@@ -158,6 +158,7 @@ public class Catalog {
                 new Functions.ST_GeometricMedian(),
                 new Functions.ST_NumPoints(),
                 new Functions.ST_Force3D(),
+                new Functions.ST_Force3DM(),
                 new Functions.ST_Force3DZ(),
                 new Functions.ST_ForceCollection(),
                 new Functions.ST_ForcePolygonCW(),
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 c1e031b83..214ba767e 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
@@ -1101,6 +1101,22 @@ public class Functions {
         }
     }
 
+    public static class ST_Force3DM 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,
+                             @DataTypeHint("Double") Double zValue) {
+            Geometry geometry = (Geometry) o;
+            return org.apache.sedona.common.Functions.force3DM(geometry, 
zValue);
+        }
+
+        @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.force3DM(geometry);
+        }
+    }
+
     public static class ST_Force3DZ 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
index 6db3f9493..a7ffb37ad 100644
--- a/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
+++ b/flink/src/test/java/org/apache/sedona/flink/FunctionTest.java
@@ -1326,6 +1326,17 @@ public class FunctionTest extends TestBase{
         assertEquals(expectedDims, actual);
     }
 
+    @Test
+    public void testForce3DM() {
+        Table geomTable = tableEnv.sqlQuery("SELECT 
ST_Force3DM(ST_GeomFromText('LINESTRING (1 2, 2 3, 3 4)')) AS geom");
+        Boolean actual = (Boolean) 
first(geomTable.select(call(Functions.ST_HasM.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(Boolean.TRUE, actual);
+
+        geomTable = tableEnv.sqlQuery("SELECT 
ST_Force3DM(ST_GeomFromText('POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0), (4 4, 4 6, 
6 6, 6 4, 4 4))')) AS geom");
+        actual = (Boolean) 
first(geomTable.select(call(Functions.ST_HasM.class.getSimpleName(), 
$("geom")))).getField(0);
+        assertEquals(Boolean.TRUE, actual);
+    }
+
     @Test
     public void testForceCollection() {
         int actual = (int) first(
diff --git a/python/sedona/sql/st_functions.py 
b/python/sedona/sql/st_functions.py
index 82ca9c2fe..efd721977 100644
--- a/python/sedona/sql/st_functions.py
+++ b/python/sedona/sql/st_functions.py
@@ -1574,6 +1574,17 @@ 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_Force3DM(geometry: ColumnOrName, mValue: Optional[Union[ColumnOrName, 
float]] = 0.0) -> Column:
+    """
+    Return a geometry with a 3D coordinate of value 'mValue' forced upon it. 
No change happens if the geometry is already 3D
+    :param mValue: Optional value of m coordinate to be potentially added, 
default value is 0.0
+    :param geometry: Geometry column to make 3D
+    :return: 3D geometry with either already present m coordinate if any, or m 
coordinate with given mValue
+    """
+    args = (geometry, mValue)
+    return _call_st_function("ST_Force3DM", args)
+
 @validate_argument_types
 def ST_Force3DZ(geometry: ColumnOrName, zValue: Optional[Union[ColumnOrName, 
float]] = 0.0) -> Column:
     """
diff --git a/python/tests/sql/test_dataframe_api.py 
b/python/tests/sql/test_dataframe_api.py
index 42818f07f..de104bc33 100644
--- a/python/tests/sql/test_dataframe_api.py
+++ b/python/tests/sql/test_dataframe_api.py
@@ -124,6 +124,7 @@ test_configurations = [
     (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_Force3DM, ("point", 1.0), "point_geom", "ST_AsText(geom)", "POINT 
M(0 1 1)"),
     (stf.ST_Force3DZ, ("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))"),
@@ -304,6 +305,7 @@ wrong_type_configurations = [
     (stf.ST_ExteriorRing, (None,)),
     (stf.ST_FlipCoordinates, (None,)),
     (stf.ST_Force_2D, (None,)),
+    (stf.ST_Force3DM, (None,)),
     (stf.ST_Force3DZ, (None,)),
     (stf.ST_ForceCollection, (None,)),
     (stf.ST_ForcePolygonCW, (None,)),
diff --git a/python/tests/sql/test_function.py 
b/python/tests/sql/test_function.py
index 7686eb81a..3ea889d64 100644
--- a/python/tests/sql/test_function.py
+++ b/python/tests/sql/test_function.py
@@ -1468,6 +1468,11 @@ class TestPredicateJoin(TestBase):
         actual = actualDf.selectExpr("ST_NDims(geom)").take(1)[0][0]
         assert expected == actual
 
+    def test_force3DM(self):
+        actualDf = self.spark.sql("SELECT 
ST_Force3DM(ST_GeomFromText('LINESTRING(0 1, 1 0, 2 0)'), 1.1) AS geom")
+        actual = actualDf.selectExpr("ST_HasM(geom)").take(1)[0][0]
+        assert actual
+
     def test_force3DZ(self):
         expected = 3
         actualDf = self.spark.sql("SELECT 
ST_Force3DZ(ST_GeomFromText('LINESTRING(0 1, 1 0, 2 0)'), 1.1) AS geom")
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 509687a6d..58a287266 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
@@ -191,6 +191,7 @@ object Catalog {
     function[ST_LengthSpheroid](),
     function[ST_NumPoints](),
     function[ST_Force3D](0.0),
+    function[ST_Force3DM](0.0),
     function[ST_Force3DZ](0.0),
     function[ST_ForceCollection](),
     function[ST_NRings](),
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 d414d30d0..d8af1de13 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
@@ -1208,6 +1208,14 @@ case class ST_Force3DZ(inputExpressions: Seq[Expression])
   }
 }
 
+case class ST_Force3DM(inputExpressions: Seq[Expression])
+  extends InferredExpression(inferrableFunction2(Functions.force3DM)) {
+
+  protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = 
{
+    copy(inputExpressions = newChildren)
+  }
+}
+
 case class ST_ForceCollection(inputExpressions: Seq[Expression])
   extends InferredExpression(Functions.forceCollection _) {
 
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 a54e59044..b8ceae4d1 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
@@ -427,6 +427,11 @@ object st_functions extends DataFrameAPI {
 
   def ST_Force3D(geometry: String, zValue: Double): Column = 
wrapExpression[ST_Force3D](geometry, zValue)
 
+  def ST_Force3DM(geometry: Column): Column = 
wrapExpression[ST_Force3DM](geometry, 0.0)
+  def ST_Force3DM(geometry: String): Column = 
wrapExpression[ST_Force3DM](geometry, 0.0)
+  def ST_Force3DM(geometry: Column, zValue: Column): Column = 
wrapExpression[ST_Force3DM](geometry, zValue)
+  def ST_Force3DM(geometry: String, zValue: Double): Column = 
wrapExpression[ST_Force3DM](geometry, zValue)
+
   def ST_Force3DZ(geometry: Column): Column = 
wrapExpression[ST_Force3DZ](geometry, 0.0)
   def ST_Force3DZ(geometry: String): Column = 
wrapExpression[ST_Force3DZ](geometry, 0.0)
   def ST_Force3DZ(geometry: Column, zValue: Column): Column = 
wrapExpression[ST_Force3DZ](geometry, zValue)
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 72a55323d..49860651e 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
@@ -1348,6 +1348,17 @@ class dataFrameAPITestScala extends TestBaseScala {
       assertEquals(expectedGeomDefaultValue, 
wktWriter.write(actualGeomDefaultValue))
     }
 
+    it("Passed ST_Force3DM") {
+      val lineDf = sparkSession.sql("SELECT ST_GeomFromWKT('LINESTRING (0 1, 1 
0, 2 0)') AS geom")
+      val expectedGeom = "LINESTRING M(0 1 2.3, 1 0 2.3, 2 0 2.3)"
+      val expectedGeomDefaultValue = "LINESTRING M(0 1 0, 1 0 0, 2 0 0)"
+      val forcedGeom = lineDf.select(ST_AsText(ST_Force3DM("geom", 
2.3))).take(1)(0).get(0).asInstanceOf[String]
+      assertEquals(expectedGeom, forcedGeom)
+      val lineDfDefaultValue = sparkSession.sql("SELECT 
ST_GeomFromWKT('LINESTRING (0 1, 1 0, 2 0)') AS geom")
+      val actualGeomDefaultValue = 
lineDfDefaultValue.select(ST_AsText(ST_Force3DM("geom"))).take(1)(0).get(0).asInstanceOf[String]
+      assertEquals(expectedGeomDefaultValue, actualGeomDefaultValue)
+    }
+
     it("Passed ST_Force3DZ") {
       val lineDf = sparkSession.sql("SELECT ST_GeomFromWKT('LINESTRING (0 1, 1 
0, 2 0)') AS geom")
       val expectedGeom = "LINESTRING Z(0 1 2.3, 1 0 2.3, 2 0 2.3)"
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 cf7e40460..0bff17b5f 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
@@ -2292,6 +2292,24 @@ class functionTestScala extends TestBaseScala with 
Matchers with GeometrySample
     }
   }
 
+  it("Should pass ST_Force3DM") {
+    val geomTestCases = Map(
+      ("'LINESTRING (0 1, 1 0, 2 0)'") -> ("'LINESTRING M(0 1 1, 1 0 1, 2 0 
1)'", "'LINESTRING M(0 1 0, 1 0 0, 2 0 0)'"),
+      ("'LINESTRING M(0 1 3, 1 0 3, 2 0 3)'") -> ("'LINESTRING M(0 1 3, 1 0 3, 
2 0 3)'", "'LINESTRING M(0 1 3, 1 0 3, 2 0 3)'"),
+      ("'LINESTRING EMPTY'") -> ("'LINESTRING EMPTY'", "'LINESTRING EMPTY'")
+    )
+    for (((geom), expectedResult) <- geomTestCases) {
+      val df = sparkSession.sql(s"SELECT 
ST_AsText(ST_Force3DM(ST_GeomFromWKT($geom), 1)) AS geom, " + 
s"$expectedResult")
+      val dfDefaultValue = sparkSession.sql(s"SELECT 
ST_AsText(ST_Force3DM(ST_GeomFromWKT($geom))) AS geom, " + s"$expectedResult")
+      val actual = df.take(1)(0).get(0).asInstanceOf[String]
+      val expected = 
df.take(1)(0).get(1).asInstanceOf[GenericRowWithSchema].get(0).asInstanceOf[String]
+      val actualDefaultValue = 
dfDefaultValue.take(1)(0).get(0).asInstanceOf[String]
+      val expectedDefaultValue = 
dfDefaultValue.take(1)(0).get(1).asInstanceOf[GenericRowWithSchema].get(1).asInstanceOf[String]
+      assertEquals(expected, actual)
+      assertEquals(expectedDefaultValue, actualDefaultValue);
+    }
+  }
+
   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)

Reply via email to