diff --git a/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DTransformer.java b/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DTransformer.java index ab8b477675..f56f03dddb 100644 --- a/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DTransformer.java +++ b/common/src/main/java/org/apache/sedona/common/utils/GeometryForce3DTransformer.java @@ -49,4 +49,16 @@ public static Geometry transform(Geometry geometry, double zValue) { GeometryForce3DTransformer transformer = new GeometryForce3DTransformer(zValue); return transformer.transform(geometry); } + + protected Geometry transformMultiPolygon(MultiPolygon geom, Geometry parent) { + // Transform each polygon individually to avoid recursion + Polygon[] transformedPolygons = new Polygon[geom.getNumGeometries()]; + for (int i = 0; i < geom.getNumGeometries(); i++) { + Polygon polygon = (Polygon) geom.getGeometryN(i); + Geometry transformed = super.transform(polygon); + transformedPolygons[i] = (Polygon) transformed; + } + // Always return as MultiPolygon, even if there's only one polygon + return geom.getFactory().createMultiPolygon(transformedPolygons); + } } 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 15bff34e23..71983f7344 100644 --- a/common/src/test/java/org/apache/sedona/common/FunctionsTest.java +++ b/common/src/test/java/org/apache/sedona/common/FunctionsTest.java @@ -2460,6 +2460,24 @@ public void force3DHybridGeomCollectionDefaultValue() { wktWriter3D.write(actualGeometryCollection.getGeometryN(0).getGeometryN(2))); } + @Test + public void force3DMultiPolygonWithSinglePolygon() { + // Test that a MultiPolygon with a single polygon remains a MultiPolygon after force3D + Polygon polygon = GEOMETRY_FACTORY.createPolygon(coordArray(0, 0, 10, 0, 10, 10, 0, 10, 0, 0)); + MultiPolygon multiPolygon = GEOMETRY_FACTORY.createMultiPolygon(new Polygon[] {polygon}); + + Geometry forced = Functions.force3D(multiPolygon, 5.0); + + assertTrue(forced instanceof MultiPolygon); + assertEquals(1, ((MultiPolygon) forced).getNumGeometries()); + + Polygon forcedPolygon = (Polygon) ((MultiPolygon) forced).getGeometryN(0); + Coordinate[] coords = forcedPolygon.getCoordinates(); + for (Coordinate coord : coords) { + assertEquals(5.0, coord.getZ(), 0.0001); + } + } + @Test public void makeLine() { Point point1 = GEOMETRY_FACTORY.createPoint(new Coordinate(0, 0)); 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 b305b9f7ae..d905cb424b 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 @@ -3033,6 +3033,17 @@ class functionTestScala } } + it("should pass ST_Force3D with MultiPolygon containing single polygon") { + // Test that a MultiPolygon with a single polygon remains a MultiPolygon after force3D + val df = sparkSession.sql( + "SELECT ST_AsText(ST_Force3D(ST_GeomFromWKT('MULTIPOLYGON (((0 0, 10 0, 10 10, 0 10, 0 0)))'), 5.0)) AS geom") + val actual = df.take(1)(0).get(0).asInstanceOf[String] + // Should still be MULTIPOLYGON, not POLYGON + assertTrue(actual.startsWith("MULTIPOLYGON")) + assertTrue(actual.contains("Z")) + assertTrue(actual.contains("5")) + } + it("Should pass ST_Force3DZ") { val geomTestCases = Map( ("'LINESTRING (0 1, 1 0, 2 0)'") -> ("'LINESTRING Z(0 1 1, 1 0 1, 2 0 1)'", "'LINESTRING Z(0 1 0, 1 0 0, 2 0 0)'"),