[GH-3369] Preserve WKB dimensions in file and adapter readers

common/src/main/java/org/apache/sedona/common/S2Geography/WKBGeography.java

@@ -130,8 +130,7 @@ public Geometry getJTSGeometry() {
         result = jtsGeometry;
         if (result == null) {
           try {
-            org.datasyslab.jts.io.WKBReader reader = new org.datasyslab.jts.io.WKBReader();
-            result = reader.read(wkbBytes);
+            result = org.datasyslab.jts.io.WKBReader.forDeclaredDimensions().read(wkbBytes);
             result.setSRID(getSRID());
           } catch (ParseException e) {
             throw new RuntimeException("Failed to parse WKB to JTS Geometry", e);

common/src/test/java/org/apache/sedona/common/S2Geography/WKBGeographyTest.java

@@ -25,6 +25,7 @@
 import com.google.common.geometry.S2Point;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.EnumSet;
 import java.util.List;
 import org.apache.sedona.common.geography.Constructors;
 import org.apache.sedona.common.geography.Functions;
@@ -35,6 +36,7 @@
 import org.locationtech.jts.geom.Point;
 import org.locationtech.jts.geom.PrecisionModel;
 import org.locationtech.jts.io.ByteOrderValues;
+import org.locationtech.jts.io.Ordinate;
 import org.locationtech.jts.io.ParseException;
 
 public class WKBGeographyTest {
@@ -66,6 +68,35 @@ public void fromWKB_point_lazyParse() throws ParseException {
     assertEquals(4326, jts.getSRID());
   }
 
+  @Test
+  public void fromWKB_emptyJtsGeometryRetainsDeclaredLayout() {
+    GeometryFactory factory = new GeometryFactory();
+    int[][] layouts = {{3, 0}, {3, 1}, {4, 1}};
+    for (int[] layout : layouts) {
+      int dimension = layout[0];
+      int measures = layout[1];
+      EnumSet<Ordinate> ordinates = EnumSet.of(Ordinate.X, Ordinate.Y);
+      if (dimension - measures > 2) ordinates.add(Ordinate.Z);
+      if (measures > 0) ordinates.add(Ordinate.M);
+      org.locationtech.jts.io.WKBWriter writer = new org.locationtech.jts.io.WKBWriter(dimension);
+      writer.setOutputOrdinates(ordinates);
+      Geometry empty =
+          factory.createPolygon(
+              factory.createLinearRing(
+                  factory.getCoordinateSequenceFactory().create(0, dimension, measures)));
+
+      Geometry parsed = WKBGeography.fromWKB(writer.write(empty), 4326).getJTSGeometry();
+      Geometry result =
+          org.apache.sedona.common.geometrySerde.GeometrySerializer.deserialize(
+              org.apache.sedona.common.geometrySerde.GeometrySerializer.serialize(parsed));
+      org.locationtech.jts.geom.CoordinateSequence sequence =
+          ((org.locationtech.jts.geom.Polygon) result).getExteriorRing().getCoordinateSequence();
+      assertEquals(dimension, sequence.getDimension());
+      assertEquals(measures, sequence.getMeasures());
+      assertEquals(4326, result.getSRID());
+    }
+  }
+
   @Test
   public void fromJTS_point() {
     GeometryFactory gf = new GeometryFactory();

docs/tutorial/files/geopackage-sedona-spark.md

@@ -79,6 +79,8 @@ Here are the contents of the DataFrame:
 
 The geometry column can contain many different geometric objects like points, polygons, and many more.
 
+Since Sedona 2.0.0, reads preserve declared Z/M dimensions even for empty points, lines, and polygons or NaN ordinates. If collecting mixed layouts raises a heterogeneous-layout error, see [ST_Collect](../../api/sql/Geometry-Editors/ST_Collect.md).
+
 You can also see the metadata of the GeoPackage file:
 
 ```python

docs/tutorial/files/geoparquet-sedona-spark.md

@@ -88,6 +88,8 @@ Here are the results:
 +---+---------------------+
 ```
 
+Since Sedona 2.0.0, reads preserve declared Z/M dimensions even for empty points, lines, and polygons or NaN ordinates. If collecting mixed layouts raises a heterogeneous-layout error, see [ST_Collect](../../api/sql/Geometry-Editors/ST_Collect.md).
+
 Here's how Sedona executes this query under the hood:
 
 1. It fetches the schema from the footer of a GeoParquet file, so no schema inference is needed.

spark/common/src/main/scala/org/apache/sedona/python/wrapper/translation/PythonGeometrySerializer.scala

@@ -20,8 +20,8 @@ package org.apache.sedona.python.wrapper.translation
 
 import org.apache.sedona.common.geometryObjects.Circle
 import org.apache.sedona.python.wrapper.SerializationException
-import org.locationtech.jts.geom.Geometry
 import org.datasyslab.jts.io.WKBReader
+import org.locationtech.jts.geom.Geometry
 
 import java.nio.ByteBuffer
 
@@ -47,13 +47,12 @@ private[python] class PythonGeometrySerializer extends Serializable {
   }
 
   def deserialize(isCircle: Int, values: Array[Byte], offset: Int): Geometry = {
-    val reader = new WKBReader()
     if (isCircle == 1) {
-      val geom = reader.read(values.slice(offset + 8, values.length))
+      val geom = WKBReader.forDeclaredDimensions().read(values.slice(offset + 8, values.length))
       val radius = ByteBuffer.wrap(values.slice(offset, offset + 8)).getDouble()
       new Circle(geom, radius)
     } else if (isCircle == 0) {
-      reader.read(values.slice(offset, values.length))
+      WKBReader.forDeclaredDimensions().read(values.slice(offset, values.length))
     } else {
       throw SerializationException("Can not deserialize object")
     }

spark/common/src/main/scala/org/apache/sedona/sql/datasources/geopackage/transform/GeometryReader.scala

@@ -20,7 +20,6 @@ package org.apache.sedona.sql.datasources.geopackage.transform
 
 import org.apache.sedona.sql.datasources.geopackage.errors.GeopackageException
 import org.apache.spark.sql.sedona_sql.UDT.GeometryUDT
-import org.locationtech.jts.geom.{GeometryFactory, PrecisionModel}
 import org.datasyslab.jts.io.WKBReader
 
 import java.nio.{ByteBuffer, ByteOrder}
@@ -56,8 +55,7 @@ object GeometryReader {
     val wkb = new Array[Byte](reader.remaining())
     reader.get(wkb)
 
-    val wkbReader = new WKBReader(new GeometryFactory(new PrecisionModel(), srid))
-    val geom = wkbReader.read(wkb)
+    val geom = WKBReader.forDeclaredDimensions(srid).read(wkb)
 
     // that needs rewriting
     GeometryUDT.serialize(geom)

spark/common/src/main/scala/org/apache/spark/sql/execution/datasources/geoparquet/GeoParquetRowConverter.scala

@@ -220,8 +220,7 @@ private[geoparquet] class GeoParquetRowConverter(
         if (parquetType.isPrimitive) {
           new ParquetPrimitiveConverter(updater) {
             override def addBinary(value: Binary): Unit = {
-              val wkbReader = new WKBReader()
-              val geom = wkbReader.read(value.getBytes)
+              val geom = WKBReader.forDeclaredDimensions().read(value.getBytes)
               geom.setSRID(srid)
               this.updater.set(GeometryUDT.serialize(geom))
             }
@@ -233,9 +232,8 @@ private[geoparquet] class GeoParquetRowConverter(
               ArrayType(ByteType, containsNull = false),
               updater) {
               override def end(): Unit = {
-                val wkbReader = new WKBReader()
                 val byteArray = currentArray.map(_.asInstanceOf[Byte]).toArray
-                val geom = wkbReader.read(byteArray)
+                val geom = WKBReader.forDeclaredDimensions().read(byteArray)
                 geom.setSRID(srid)
                 this.updater.set(GeometryUDT.serialize(geom))
               }

spark/common/src/test/scala/org/apache/sedona/python/wrapper/TestToPythonSerialization.scala

@@ -25,11 +25,12 @@ import org.apache.sedona.python.wrapper.utils.implicits._
 import org.apache.sedona.sql.TestBaseScala
 import org.apache.spark.api.java.JavaPairRDD
 import org.locationtech.jts.geom.{Geometry, GeometryFactory}
-import org.locationtech.jts.io.WKTReader
+import org.locationtech.jts.io.{Ordinate, WKBWriter, WKTReader}
 import org.scalatest.matchers.must.Matchers.contain
 import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
 
 import java.io.{FileInputStream, InputStream}
+import java.util.EnumSet
 import scala.io.Source
 import scala.jdk.CollectionConverters._
 
@@ -95,6 +96,28 @@ class TestToPythonSerialization extends TestBaseScala {
         .flatMap(samplePoint => pythonGeometrySerializer.serialize(samplePoint)))
 
   describe("Sedona Python Wrapper Test") {
+    it("retains the declared layout of inbound empty WKB") {
+      Seq((3, 0), (3, 1), (4, 1)).foreach { case (dimension, measures) =>
+        val sequence = geometryFactory.getCoordinateSequenceFactory.create(0, dimension, measures)
+        val empty = geometryFactory.createPolygon(geometryFactory.createLinearRing(sequence))
+        val writer = new WKBWriter(dimension)
+        val ordinates = EnumSet.of(Ordinate.X, Ordinate.Y)
+        if (dimension - measures > 2) ordinates.add(Ordinate.Z)
+        if (measures > 0) ordinates.add(Ordinate.M)
+        writer.setOutputOrdinates(ordinates)
+
+        val parsed = pythonGeometrySerializer.deserialize(0, writer.write(empty), 0)
+        val result = org.apache.sedona.common.geometrySerde.GeometrySerializer.deserialize(
+          org.apache.sedona.common.geometrySerde.GeometrySerializer.serialize(parsed))
+        val resultSequence = result
+          .asInstanceOf[org.locationtech.jts.geom.Polygon]
+          .getExteriorRing
+          .getCoordinateSequence
+        resultSequence.getDimension shouldBe dimension
+        resultSequence.getMeasures shouldBe measures
+      }
+    }
+
     it("Test Serialize To Python JavaRDD[Geometry]") {
       val convertedToPythonRDD =
         GeometryRddConverter(pointSpatialRDD, pythonGeometrySerializer).translateToPython

spark/common/src/test/scala/org/apache/sedona/sql/datasources/geopackage/transform/GeometryReaderTest.scala

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.sql.datasources.geopackage.transform
+
+import org.apache.spark.sql.sedona_sql.UDT.GeometryUDT
+import org.locationtech.jts.geom.{GeometryFactory, Polygon}
+import org.locationtech.jts.io.{Ordinate, WKBWriter}
+import org.scalatest.funsuite.AnyFunSuite
+
+import java.nio.{ByteBuffer, ByteOrder}
+import java.util.EnumSet
+
+class GeometryReaderTest extends AnyFunSuite {
+  test("extractWKB retains an empty geometry's declared layout and GeoPackage SRID") {
+    val factory = new GeometryFactory()
+    Seq((3, 0), (3, 1), (4, 1)).foreach { case (dimension, measures) =>
+      val sequence = factory.getCoordinateSequenceFactory.create(0, dimension, measures)
+      val empty = factory.createPolygon(factory.createLinearRing(sequence))
+      val writer = new WKBWriter(dimension)
+      val ordinates = EnumSet.of(Ordinate.X, Ordinate.Y)
+      if (dimension - measures > 2) ordinates.add(Ordinate.Z)
+      if (measures > 0) ordinates.add(Ordinate.M)
+      writer.setOutputOrdinates(ordinates)
+      val wkb = writer.write(empty)
+      val bytes = ByteBuffer
+        .allocate(8 + wkb.length)
+        .order(ByteOrder.LITTLE_ENDIAN)
+        .put('G'.toByte)
+        .put('P'.toByte)
+        .put(0.toByte)
+        .put(0x11.toByte)
+        .putInt(4326)
+        .put(wkb)
+        .array()
+
+      val result = GeometryUDT.deserialize(GeometryReader.extractWKB(bytes)).asInstanceOf[Polygon]
+      val resultSequence = result.getExteriorRing.getCoordinateSequence
+      assert(resultSequence.getDimension == dimension)
+      assert(resultSequence.getMeasures == measures)
+      assert(result.getSRID == 4326)
+    }
+  }
+}

spark/common/src/test/scala/org/apache/spark/sql/execution/datasources/geoparquet/GeoParquetRowConverterTest.scala

@@ -0,0 +1,151 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.spark.sql.execution.datasources.geoparquet
+
+import org.apache.parquet.io.api.Binary
+import org.apache.parquet.io.api.GroupConverter
+import org.apache.parquet.schema.MessageTypeParser
+import org.apache.spark.sql.execution.datasources.geoparquet.internal.{LegacyBehaviorPolicy, NoopUpdater, RebaseSpec}
+import org.apache.spark.sql.sedona_sql.UDT.GeometryUDT
+import org.apache.spark.sql.types.{StructField, StructType}
+import org.locationtech.jts.geom.{Coordinate, GeometryFactory, Point}
+import org.locationtech.jts.io.{Ordinate, WKBWriter}
+import org.scalatest.funsuite.AnyFunSuite
+
+import java.util.{Collections, EnumSet}
+
+class GeoParquetRowConverterTest extends AnyFunSuite {
+  private def primitiveConverter(): GeoParquetRowConverter = {
+    val parquetSchema =
+      MessageTypeParser.parseMessageType("message root { optional binary geometry; }")
+    val catalystSchema = StructType(Seq(StructField("geometry", GeometryUDT(), nullable = true)))
+    val metadata = Collections.singletonMap(
+      "geo",
+      """{"version":"1.1.0","primary_column":"geometry","columns":{"geometry":{"encoding":"WKB","geometry_types":[],"crs":null}}}""")
+    new GeoParquetRowConverter(
+      new GeoParquetToSparkSchemaConverter(metadata, parameters = Map.empty),
+      parquetSchema,
+      catalystSchema,
+      None,
+      RebaseSpec(LegacyBehaviorPolicy.CORRECTED),
+      RebaseSpec(LegacyBehaviorPolicy.CORRECTED),
+      Map.empty,
+      NoopUpdater)
+  }
+
+  test("primitive WKB retains declared empty coordinate layouts through GeometryUDT") {
+    val converter = primitiveConverter()
+    val factory = new GeometryFactory()
+
+    Seq((3, 0), (3, 1), (4, 1)).foreach { case (dimension, measures) =>
+      val sequence = factory.getCoordinateSequenceFactory.create(0, dimension, measures)
+      val empty = factory.createPolygon(factory.createLinearRing(sequence))
+      val writer = new WKBWriter(dimension)
+      val ordinates = EnumSet.of(Ordinate.X, Ordinate.Y)
+      if (dimension - measures > 2) ordinates.add(Ordinate.Z)
+      if (measures > 0) ordinates.add(Ordinate.M)
+      writer.setOutputOrdinates(ordinates)
+
+      converter.start()
+      converter
+        .getConverter(0)
+        .asPrimitiveConverter()
+        .addBinary(Binary.fromConstantByteArray(writer.write(empty)))
+      converter.end()
+      val result = GeometryUDT.deserialize(converter.currentRecord.getBinary(0))
+      val resultSequence = result
+        .asInstanceOf[org.locationtech.jts.geom.Polygon]
+        .getExteriorRing
+        .getCoordinateSequence
+      assert(resultSequence.getDimension == dimension)
+      assert(resultSequence.getMeasures == measures)
+    }
+  }
+
+  test("primitive WKB retains NaN Z without carrying its layout into later XY rows") {
+    val converter = primitiveConverter()
+    val factory = new GeometryFactory()
+    val nanZ = new WKBWriter(3).write(factory.createPoint(new Coordinate(1, 2, 9)))
+    java.nio.ByteBuffer.wrap(nanZ).putDouble(21, Double.NaN)
+    val xy = new WKBWriter(2).write(factory.createPoint(new Coordinate(1, 2)))
+
+    Seq((nanZ, 3), (xy, 2), (nanZ, 3)).foreach { case (wkb, dimension) =>
+      converter.start()
+      converter
+        .getConverter(0)
+        .asPrimitiveConverter()
+        .addBinary(Binary.fromConstantByteArray(wkb))
+      converter.end()
+      val result =
+        GeometryUDT.deserialize(converter.currentRecord.getBinary(0)).asInstanceOf[Point]
+      val sequence = result.getCoordinateSequence
+      assert(!result.isEmpty)
+      assert(sequence.getDimension == dimension)
+      assert(sequence.getMeasures == 0)
+      assert(sequence.getX(0) == 1)
+      assert(sequence.getY(0) == 2)
+      if (dimension == 3) assert(sequence.getZ(0).isNaN)
+    }
+  }
+
+  test("legacy byte-array WKB retains declared empty coordinate layouts through GeometryUDT") {
+    val parquetSchema = MessageTypeParser.parseMessageType(
+      "message root { optional group geometry (LIST) { repeated int32 array (INT_8); } }")
+    val catalystSchema = StructType(Seq(StructField("geometry", GeometryUDT(), nullable = true)))
+    val metadata = Collections.singletonMap(
+      "geo",
+      """{"version":"1.1.0","primary_column":"geometry","columns":{"geometry":{"encoding":"WKB","geometry_types":[],"crs":null}}}""")
+    val parameters = Map("legacyMode" -> "true")
+    val converter = new GeoParquetRowConverter(
+      new GeoParquetToSparkSchemaConverter(metadata, parameters = parameters),
+      parquetSchema,
+      catalystSchema,
+      None,
+      RebaseSpec(LegacyBehaviorPolicy.CORRECTED),
+      RebaseSpec(LegacyBehaviorPolicy.CORRECTED),
+      parameters,
+      NoopUpdater)
+    val factory = new GeometryFactory()
+
+    Seq((3, 0), (3, 1), (4, 1)).foreach { case (dimension, measures) =>
+      val sequence = factory.getCoordinateSequenceFactory.create(0, dimension, measures)
+      val empty = factory.createPolygon(factory.createLinearRing(sequence))
+      val writer = new WKBWriter(dimension)
+      val ordinates = EnumSet.of(Ordinate.X, Ordinate.Y)
+      if (dimension - measures > 2) ordinates.add(Ordinate.Z)
+      if (measures > 0) ordinates.add(Ordinate.M)
+      writer.setOutputOrdinates(ordinates)
+
+      converter.start()
+      val geometryConverter = converter.getConverter(0).asInstanceOf[GroupConverter]
+      geometryConverter.start()
+      val byteConverter = geometryConverter.getConverter(0).asPrimitiveConverter()
+      writer.write(empty).foreach(byte => byteConverter.addInt(byte))
+      geometryConverter.end()
+      converter.end()
+      val result = GeometryUDT.deserialize(converter.currentRecord.getBinary(0))
+      val resultSequence = result
+        .asInstanceOf[org.locationtech.jts.geom.Polygon]
+        .getExteriorRing
+        .getCoordinateSequence
+      assert(resultSequence.getDimension == dimension)
+      assert(resultSequence.getMeasures == measures)
+    }
+  }
+}