[GH-2504] Geopandas: Implement force_3d (#2512)

chay0112 · web-flow · commit c7e42cbbfeec · 2025-12-03T21:51:25.000-08:00
diff --git a/python/sedona/spark/geopandas/base.py b/python/sedona/spark/geopandas/base.py
@@ -949,8 +949,70 @@ def force_2d(self):
         """
         return _delegate_to_geometry_column("force_2d", self)
 
-    # def force_3d(self, z=0):
-    #     raise NotImplementedError("This method is not implemented yet.")
+    def force_3d(self, z=0.0):
+        """Force the dimensionality of a geometry to 3D.
+
+        2D geometries will get the provided Z coordinate; 3D geometries
+        are unchanged (unless their Z coordinate is ``np.nan``).
+
+        Note: Sedona's behavior may differ from Geopandas' for M and ZM geometries.
+        For M geometries, Sedona will replace the M coordinate and add the Z coordinate.
+        For ZM geometries, Sedona will drop the M coordinate and retain the Z coordinate.
+
+        Parameters
+        ----------
+        z : float | array_like (default 0)
+            Z coordinate to be assigned
+
+        Returns
+        -------
+        GeoSeries
+
+        Examples
+        --------
+        >>> from shapely import Polygon, LineString, Point
+        >>> from sedona.spark.geopandas import GeoSeries
+        >>> s = GeoSeries(
+        ...     [
+        ...         Point(1, 2),
+        ...         Point(0.5, 2.5, 2),
+        ...         LineString([(1, 1), (0, 1), (1, 0)]),
+        ...         Polygon([(0, 0), (0, 10), (10, 10)]),
+        ...     ],
+        ... )
+        >>> s
+        0                          POINT (1 2)
+        1                  POINT Z (0.5 2.5 2)
+        2           LINESTRING (1 1, 0 1, 1 0)
+        3    POLYGON ((0 0, 0 10, 10 10, 0 0))
+        dtype: geometry
+
+        >>> s.force_3d()
+        0                                POINT Z (1 2 0)
+        1                            POINT Z (0.5 2.5 2)
+        2             LINESTRING Z (1 1 0, 0 1 0, 1 0 0)
+        3    POLYGON Z ((0 0 0, 0 10 0, 10 10 0, 0 0 0))
+        dtype: geometry
+
+        Z coordinate can be specified as scalar:
+
+        >>> s.force_3d(4)
+        0                                POINT Z (1 2 4)
+        1                            POINT Z (0.5 2.5 2)
+        2             LINESTRING Z (1 1 4, 0 1 4, 1 0 4)
+        3    POLYGON Z ((0 0 4, 0 10 4, 10 10 4, 0 0 4))
+        dtype: geometry
+
+        Or as an array-like (one value per geometry):
+
+        >>> s.force_3d(range(4))
+        0                                POINT Z (1 2 0)
+        1                            POINT Z (0.5 2.5 2)
+        2             LINESTRING Z (1 1 2, 0 1 2, 1 0 2)
+        3    POLYGON Z ((0 0 3, 0 10 3, 10 10 3, 0 0 3))
+        dtype: geometry
+        """
+        return _delegate_to_geometry_column("force_3d", self, z)
 
     # def line_merge(self, directed=False):
     #     raise NotImplementedError("This method is not implemented yet.")
diff --git a/python/sedona/spark/geopandas/geoseries.py b/python/sedona/spark/geopandas/geoseries.py
@@ -1093,9 +1093,17 @@ def force_2d(self) -> "GeoSeries":
         spark_expr = stf.ST_Force_2D(self.spark.column)
         return self._query_geometry_column(spark_expr, returns_geom=True)
 
-    def force_3d(self, z=0):
-        # Implementation of the abstract method.
-        raise NotImplementedError("This method is not implemented yet.")
+    def force_3d(self, z=0.0) -> "GeoSeries":
+        other_series, extended = self._make_series_of_val(z)
+        align = not extended
+
+        spark_expr = stf.ST_Force3D(F.col("L"), F.col("R"))
+        return self._row_wise_operation(
+            spark_expr,
+            other_series,
+            align=align,
+            returns_geom=True,
+        )
 
     def line_merge(self, directed=False):
         # Implementation of the abstract method.
diff --git a/python/tests/geopandas/test_geoseries.py b/python/tests/geopandas/test_geoseries.py
@@ -1520,7 +1520,95 @@ def test_force_2d(self):
         self.check_sgpd_equals_gpd(df_result, expected)
 
     def test_force_3d(self):
-        pass
+        # 1. 2D geometries promoted to 3D with default z=0.0
+        s = sgpd.GeoSeries(
+            [
+                Point(1, 2),
+                Point(0.5, 2.5, 2),
+                Point(1, 1, np.nan),
+                LineString([(1, 1), (0, 1), (1, 0)]),
+                Polygon([(0, 0), (0, 10), (10, 10)]),
+                GeometryCollection(
+                    [
+                        Point(1, 1),
+                        LineString([(1, 1), (0, 1), (1, 0)]),
+                    ]
+                ),
+            ]
+        )
+        # Promote 2D to 3D with z=0, keep 3D as is
+        expected = gpd.GeoSeries(
+            [
+                Point(1, 2, 0),
+                Point(0.5, 2.5, 2),
+                Point(1, 1, 0),
+                LineString([(1, 1, 0), (0, 1, 0), (1, 0, 0)]),
+                Polygon([(0, 0, 0), (0, 10, 0), (10, 10, 0), (0, 0, 0)]),
+                GeometryCollection(
+                    [
+                        Point(1, 1, 0),
+                        LineString([(1, 1, 0), (0, 1, 0), (1, 0, 0)]),
+                    ]
+                ),
+            ]
+        )
+        result = s.force_3d()
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # 2. 2D geometries promoted to 3D with scalar z
+        expected = gpd.GeoSeries(
+            [
+                Point(1, 2, 4),
+                Point(0.5, 2.5, 2),
+                Point(1, 1, 4),
+                LineString([(1, 1, 4), (0, 1, 4), (1, 0, 4)]),
+                Polygon([(0, 0, 4), (0, 10, 4), (10, 10, 4), (0, 0, 4)]),
+                GeometryCollection(
+                    [
+                        Point(1, 1, 4),
+                        LineString([(1, 1, 4), (0, 1, 4), (1, 0, 4)]),
+                    ]
+                ),
+            ]
+        )
+        result = s.force_3d(4)
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # 3. Array-like z: use ps.Series
+        z = [0, 2, 2, 3, 4, 5]
+        expected = gpd.GeoSeries(
+            [
+                Point(1, 2, 0),
+                Point(0.5, 2.5, 2),
+                Point(1, 1, 2),
+                LineString([(1, 1, 3), (0, 1, 3), (1, 0, 3)]),
+                Polygon([(0, 0, 4), (0, 10, 4), (10, 10, 4), (0, 0, 4)]),
+                GeometryCollection(
+                    [
+                        Point(1, 1, 5),
+                        LineString([(1, 1, 5), (0, 1, 5), (1, 0, 5)]),
+                    ]
+                ),
+            ]
+        )
+        result = s.force_3d(z)
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # 4. Ensure M and ZM geometries are handled correctly
+        s = sgpd.GeoSeries(
+            [
+                shapely.wkt.loads("POINT M (1 2 3)"),
+                shapely.wkt.loads("POINT ZM (1 2 3 4)"),
+            ]
+        )
+        result = s.force_3d(7.5)
+        expected = gpd.GeoSeries(
+            [
+                shapely.wkt.loads("POINT Z (1 2 7.5)"),
+                shapely.wkt.loads("POINT Z (1 2 3)"),
+            ]
+        )
+        self.check_sgpd_equals_gpd(result, expected)
 
     def test_line_merge(self):
         pass
diff --git a/python/tests/geopandas/test_match_geopandas_series.py b/python/tests/geopandas/test_match_geopandas_series.py
@@ -878,7 +878,48 @@ def test_force_2d(self):
         self.check_sgpd_equals_gpd(sgpd_3d, gpd_3d)
 
     def test_force_3d(self):
-        pass
+        # force_3d was added from geopandas 1.0.0
+        if parse_version(gpd.__version__) < parse_version("1.0.0"):
+            pytest.skip("geopandas force_3d requires version 1.0.0 or higher")
+        # 1) Promote 2D to 3D with z = 4
+        for geom in self.geoms:
+            if isinstance(geom[0], (LinearRing)):
+                continue
+            sgpd_result = GeoSeries(geom).force_3d(4)
+            gpd_result = gpd.GeoSeries(geom).force_3d(4)
+            self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
+
+        # 2) Minimal sample for various geometry types with custom z=7.5
+        data = [
+            Point(1, 2),  # 2D
+            Point(0.5, 2.5, 2),  # 3D (Z)
+            LineString([(1, 1), (0, 1), (1, 0)]),  # 2D
+            Polygon([(0, 0), (0, 10), (10, 10)]),  # 2D
+        ]
+        sgpd_result = GeoSeries(data).force_3d(7.5)
+        gpd_result = gpd.GeoSeries(data).force_3d(7.5)
+        self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
+
+        # 3) Array-like z tests
+        geoms = self.polygons
+        lst = list(range(1, len(geoms) + 1))
+
+        # Traditional python list
+        sgpd_result = GeoSeries(geoms).force_3d(lst)
+        gpd_result = gpd.GeoSeries(geoms).force_3d(lst)
+        self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
+
+        # numpy array
+        np_array = np.array(lst)
+        sgpd_result = GeoSeries(geoms).force_3d(np_array)
+        gpd_result = gpd.GeoSeries(geoms).force_3d(np_array)
+        self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
+
+        # pandas-on-Spark Series
+        psser = ps.Series(lst)
+        sgpd_result = GeoSeries(geoms).force_3d(psser)
+        gpd_result = gpd.GeoSeries(geoms).force_3d(psser.to_pandas())
+        self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
 
     def test_line_merge(self):
         pass
