[GH-2525] bug: Make ST_Force3D return multipolygons even if it's a single polygon inside (#2526)

Author: petern48

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);
+  }
 }

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));

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)'"),