[GH-3017] Implement rotate (#3018)

Author: oglego

python/sedona/spark/geopandas/base.py

@@ -1362,6 +1362,57 @@ def segmentize(self, max_segment_length):
     # def transform(self, transformation, include_z=False):
     #     raise NotImplementedError("This method is not implemented yet.")
 
+    def rotate(self, angle, origin="center", use_radians=False):
+        """Return a ``GeoSeries`` with rotated geometries.
+
+        See http://shapely.readthedocs.io/en/latest/manual.html#shapely.affinity.rotate
+        for details.
+
+        Parameters
+        ----------
+        angle : float
+            The angle of rotation can be specified in either degrees (default)
+            or radians by setting use_radians=True. Positive angles are
+            counter-clockwise and negative are clockwise rotations.
+        origin : string, Point, or tuple (x, y)
+            The point of origin can be a keyword 'center' for the bounding box
+            center (default), 'centroid' for the geometry's centroid, a Point
+            object or a coordinate tuple (x, y).
+        use_radians : boolean
+            Whether to interpret the angle of rotation as degrees or radians
+
+        Examples
+        --------
+        >>> from shapely.geometry import Point, LineString, Polygon
+        >>> from sedona.spark.geopandas import GeoSeries
+        >>> s = GeoSeries(
+        ...     [
+        ...         Point(1, 1),
+        ...         LineString([(1, -1), (1, 0)]),
+        ...         Polygon([(3, -1), (4, 0), (3, 1)]),
+        ...     ]
+        ... )
+        >>> s
+        0                         POINT (1 1)
+        1              LINESTRING (1 -1, 1 0)
+        2    POLYGON ((3 -1, 4 0, 3 1, 3 -1))
+        dtype: geometry
+
+        >>> s.rotate(90)
+        0                                          POINT (1 1)
+        1                      LINESTRING (1.5 -0.5, 0.5 -0.5)
+        2    POLYGON ((4.5 -0.5, 3.5 0.5, 2.5 -0.5, 4.5 -0.5))
+        dtype: geometry
+
+        >>> s.rotate(90, origin=(0, 0))
+        0                       POINT (-1 1)
+        1              LINESTRING (1 1, 0 1)
+        2    POLYGON ((1 3, 0 4, -1 3, 1 3))
+        dtype: geometry
+
+        """
+        return _delegate_to_geometry_column("rotate", self, angle, origin, use_radians)
+
     def force_2d(self):
         """Force the dimensionality of a geometry to 2D.
 

python/sedona/spark/geopandas/geoseries.py

@@ -1160,6 +1160,36 @@ def transform(self, transformation, include_z=False):
         # Implementation of the abstract method.
         raise NotImplementedError("This method is not implemented yet.")
 
+    def rotate(self, angle, origin="center", use_radians=False) -> "GeoSeries":
+        import math
+
+        if not use_radians:
+            angle = angle * math.pi / 180.0
+        if isinstance(origin, str):
+            if origin == "center":
+                origin_x = stf.ST_X(stf.ST_Centroid(stf.ST_Envelope(self.spark.column)))
+                origin_y = stf.ST_Y(stf.ST_Centroid(stf.ST_Envelope(self.spark.column)))
+                spark_expr = stf.ST_Rotate(self.spark.column, angle, origin_x, origin_y)
+            elif origin == "centroid":
+                origin_x = stf.ST_X(stf.ST_Centroid(self.spark.column))
+                origin_y = stf.ST_Y(stf.ST_Centroid(self.spark.column))
+                spark_expr = stf.ST_Rotate(self.spark.column, angle, origin_x, origin_y)
+            else:
+                raise ValueError(
+                    f"origin must be 'center', 'centroid', a Point, or (x, y) tuple, got {origin!r}"
+                )
+        elif isinstance(origin, (tuple, list)) and len(origin) == 2:
+            spark_expr = stf.ST_Rotate(
+                self.spark.column, angle, float(origin[0]), float(origin[1])
+            )
+        elif isinstance(origin, shapely.geometry.Point):
+            spark_expr = stf.ST_Rotate(self.spark.column, angle, origin.x, origin.y)
+        else:
+            raise TypeError(
+                "origin must be 'center', 'centroid', a Point, or a two-element tuple/list"
+            )
+        return self._query_geometry_column(spark_expr, returns_geom=True)
+
     def force_2d(self) -> "GeoSeries":
         spark_expr = stf.ST_Force_2D(self.spark.column)
         return self._query_geometry_column(spark_expr, returns_geom=True)

python/tests/geopandas/test_geoseries.py

@@ -1847,6 +1847,77 @@ def test_segmentize(self):
     def test_transform(self):
         pass
 
+    def test_rotate(self):
+        geoms = [
+            Point(1, 1),
+            LineString([(1, -1), (1, 0)]),
+            Polygon([(3, -1), (4, 0), (3, 1)]),
+            None,
+        ]
+        s = GeoSeries(geoms)
+
+        # Test default
+        result = s.rotate(90)
+        expected = gpd.GeoSeries(
+            [
+                Point(1.0, 1.0),
+                LineString([(1.5, -0.5), (0.5, -0.5)]),
+                Polygon([(4.5, -0.5), (3.5, 0.5), (2.5, -0.5), (4.5, -0.5)]),
+                None,
+            ]
+        )
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # Test with explicit origin tuple (90 degrees around 0,0)
+        result = s.rotate(90, origin=(0, 0))
+        expected = gpd.GeoSeries(
+            [
+                Point(-1.0, 1.0),
+                LineString([(1.0, 1.0), (0.0, 1.0)]),
+                Polygon([(1.0, 3.0), (0.0, 4.0), (-1.0, 3.0), (1.0, 3.0)]),
+                None,
+            ]
+        )
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # Test use_radians
+        import math
+
+        result = s.rotate(math.pi / 2, origin=(0, 0), use_radians=True)
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # Test with a Point object as the origin
+        result = s.rotate(90, origin=Point(0, 0))
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # Test origin='centroid'
+        result = s.rotate(45, origin="centroid")
+        expected = gpd.GeoSeries(
+            [
+                Point(1.0, 1.0),
+                LineString(
+                    [
+                        (1.3535533905932737, -0.8535533905932737),
+                        (0.6464466094067263, -0.14644660940672627),
+                    ]
+                ),
+                Polygon(
+                    [
+                        (3.8047378541243653, -0.9428090415820636),
+                        (3.8047378541243653, 0.4714045207910314),
+                        (2.3905242917512703, 0.4714045207910314),
+                        (3.8047378541243653, -0.9428090415820636),
+                    ]
+                ),
+                None,
+            ]
+        )
+        self.check_sgpd_equals_gpd(result, expected)
+
+        # Test invalid origin strings
+        with pytest.raises((ValueError, TypeError)):
+            s.rotate(90, origin="invalid")
+
     def test_force_2d(self):
         s = sgpd.GeoSeries(
             [

python/tests/geopandas/test_match_geopandas_series.py

@@ -956,6 +956,28 @@ def test_segmentize(self):
     def test_transform(self):
         pass
 
+    @pytest.mark.parametrize("angle", [0, 45, 90, 180])
+    @pytest.mark.parametrize(
+        "origin_kwargs",
+        [
+            {},
+            {"origin": (1, 2)},
+            {"origin": [1, 2]},
+            {"origin": "centroid"},
+            {"origin": Point(2, 4)},
+        ],
+    )
+    def test_rotate(self, angle, origin_kwargs):
+        for geom in self.geoms:
+            # Filter empty geometries within each group
+            non_empty = [g for g in geom if g is not None and not g.is_empty]
+            if not non_empty:
+                continue
+
+            sgpd_result = GeoSeries(non_empty).rotate(angle, **origin_kwargs)
+            gpd_result = gpd.GeoSeries(non_empty).rotate(angle, **origin_kwargs)
+            self.check_sgpd_equals_gpd(sgpd_result, gpd_result)
+
     def test_force_2d(self):
         # force_2d was added from geopandas 1.0.0
         if parse_version(gpd.__version__) < parse_version("1.0.0"):