[GH-3260] Preserve M and ZM in Python GeometryUDT serialization

python/sedona/spark/utils/geometry_serde_general.py

@@ -21,6 +21,7 @@
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
+import shapely
 from shapely.geometry import (
     GeometryCollection,
     LinearRing,
@@ -183,6 +184,17 @@ def serialize(geom: BaseGeometry) -> Optional[Union[bytes, bytearray]]:
     if geom is None:
         return None
 
+    # Shapely 2.1 exposes ``has_m`` even when linked against GEOS < 3.12,
+    # but evaluating that property then raises for every geometry. Avoid the
+    # property entirely when the linked GEOS cannot represent measures.
+    if getattr(shapely, "geos_version", (0, 0, 0)) >= (3, 12, 0) and getattr(
+        geom, "has_m", False
+    ):
+        raise ValueError(
+            "Serializing M or ZM geometries requires geomserde_speedup "
+            "with GEOS 3.12 or newer"
+        )
+
     if isinstance(geom, Point):
         return serialize_point(geom)
     elif isinstance(geom, LineString):
@@ -214,6 +226,11 @@ def deserialize(buffer: bytes) -> Optional[BaseGeometry]:
     preamble_byte = buffer[0]
     geom_type = (preamble_byte >> 4) & 0x0F
     coord_type = (preamble_byte >> 1) & 0x07
+    if coord_type in (CoordinateType.XYM, CoordinateType.XYZM):
+        raise ValueError(
+            "Deserializing M or ZM geometries requires geomserde_speedup "
+            "with GEOS 3.12 or newer"
+        )
     num_coords = struct.unpack_from("i", buffer, 4)[0]
     if num_coords > len(buffer):
         raise ValueError("num_coords cannot be larger than buffer size")

python/src/geom_buf.c

@@ -56,8 +56,19 @@ SedonaErrorCode get_coord_seq_info_from_geom(
   if (dims == 0) {
     return SEDONA_GEOS_ERROR;
   }
-  int has_z = (dims >= 3);
-  int has_m = 0; /* libgeos does not support M dimension for now */
+  char has_z_result = dyn_GEOSHasZ_r(handle, geom);
+  char has_m_result = dyn_GEOSHasM_r == NULL ? 0 : dyn_GEOSHasM_r(handle, geom);
+  if (has_z_result == 2 || has_m_result == 2) {
+    return SEDONA_GEOS_ERROR;
+  }
+  int has_m = has_m_result != 0;
+  /*
+   * GEOSHasZ_r may report false when the first Z ordinate is NaN on older
+   * supported GEOS versions. Coordinate dimension still records the layout,
+   * and GEOSHasM_r disambiguates XYZ from XYM when measures are supported.
+   */
+  int has_z = has_z_result != 0 || dims == 4 || (dims == 3 && !has_m);
+  dims = 2 + has_z + has_m;
   CoordinateType coord_type = coordinate_type_of(has_z, has_m);
   unsigned int bytes_per_coord = get_bytes_per_coordinate(coord_type);
   int num_coords = dyn_GEOSGetNumCoordinates_r(handle, geom);

python/src/geos_c_dyn.c

@@ -156,6 +156,11 @@ int load_geos_c_from_handle(void *handle, char *err_msg, int len) {
   dyn_GEOSCoordSeq_copyToBuffer_r =
       try_load_geos_c_symbol(handle, "GEOSCoordSeq_copyToBuffer_r");
 
+  /* GEOSHasM_r was added in GEOS 3.12. Keep it optional so the extension
+   * remains compatible with the older GEOS versions bundled by supported
+   * Shapely releases. */
+  dyn_GEOSHasM_r = try_load_geos_c_symbol(handle, "GEOSHasM_r");
+
   /* Deliberately load GEOS_init_r after all other functions, so that we can
    * check if all functions were loaded by checking if GEOS_init_r was
    * loaded. */

python/src/geos_c_dyn_funcs.h

@@ -34,6 +34,9 @@ GEOS_FP_QUALIFIER int (*dyn_GEOSGeomTypeId_r)(GEOSContextHandle_t handle,
 GEOS_FP_QUALIFIER char (*dyn_GEOSHasZ_r)(GEOSContextHandle_t handle,
                                          const GEOSGeometry *g);
 
+GEOS_FP_QUALIFIER char (*dyn_GEOSHasM_r)(GEOSContextHandle_t handle,
+                                         const GEOSGeometry *g);
+
 GEOS_FP_QUALIFIER int (*dyn_GEOSGetSRID_r)(GEOSContextHandle_t handle,
                                            const GEOSGeometry *g);
 

python/tests/utils/test_geomserde_speedup.py

@@ -17,6 +17,7 @@
 
 import pytest
 import shapely
+from packaging.version import parse as parse_version
 from shapely.geometry import (
     GeometryCollection,
     LineString,
@@ -30,6 +31,7 @@
 from shapely.wkt import loads as wkt_loads
 
 from sedona.spark.utils import geometry_serde
+from sedona.spark.utils import geometry_serde_general
 
 
 class TestGeomSerdeSpeedup:
@@ -49,6 +51,14 @@ def test_linestring(self):
         ]
         self._test_serde_roundtrip(linestrings)
 
+    def test_nan_first_z_serialization_keeps_dimension(self):
+        geometry = wkt_loads("LINESTRING Z (0 0 NaN, 1 1 2)")
+
+        buffer = geometry_serde.serialize(geometry)
+        coordinate_type = (buffer[0] & 0x0F) >> 1
+
+        assert coordinate_type == geometry_serde_general.CoordinateType.XYZ
+
     def test_multi_point(self):
         multi_points = [
             wkt_loads("MULTIPOINT EMPTY"),
@@ -146,6 +156,77 @@ def test_srid_roundtrip(self):
         point2 = TestGeomSerdeSpeedup.serde_roundtrip(point)
         assert shapely.get_srid(point2) == 1000
 
+    @pytest.mark.skipif(
+        parse_version(shapely.__version__) < parse_version("2.1")
+        or getattr(shapely, "geos_version", (0, 0, 0)) < (3, 12, 0),
+        reason="M coordinates require Shapely 2.1 and GEOS 3.12 or newer",
+    )
+    @pytest.mark.parametrize(
+        "wkt",
+        [
+            "POINT M (1 2 3)",
+            "POINT ZM (1 2 3 4)",
+            "LINESTRING M (0 0 1, 2 3 4)",
+            "LINESTRING ZM (0 0 1 2, 3 4 5 6)",
+            "POLYGON M ((0 0 1, 2 0 2, 0 2 3, 0 0 1))",
+            "GEOMETRYCOLLECTION ZM (POINT ZM (1 2 3 4), "
+            "LINESTRING ZM (0 0 1 2, 3 4 5 6))",
+        ],
+    )
+    def test_m_roundtrip(self, wkt):
+        geometry = shapely.from_wkt(wkt)
+        actual = TestGeomSerdeSpeedup.serde_roundtrip(geometry)
+
+        assert shapely.to_wkt(actual) == shapely.to_wkt(geometry)
+        assert actual.has_z == geometry.has_z
+        assert actual.has_m == geometry.has_m
+
+    @pytest.mark.skipif(
+        parse_version(shapely.__version__) < parse_version("2.1")
+        or getattr(shapely, "geos_version", (0, 0, 0)) < (3, 12, 0),
+        reason="M coordinates require Shapely 2.1 and GEOS 3.12 or newer",
+    )
+    @pytest.mark.parametrize("wkt", ["POINT M (1 2 3)", "POINT ZM (1 2 3 4)"])
+    def test_general_serializer_rejects_m_instead_of_losing_it(self, wkt):
+        geometry = shapely.from_wkt(wkt)
+
+        with pytest.raises(ValueError, match="requires geomserde_speedup"):
+            geometry_serde_general.serialize(geometry)
+
+    @pytest.mark.parametrize(
+        "coord_type",
+        [
+            geometry_serde_general.CoordinateType.XYM,
+            geometry_serde_general.CoordinateType.XYZM,
+        ],
+    )
+    def test_general_deserializer_rejects_m_instead_of_losing_it(self, coord_type):
+        buffer = geometry_serde_general.create_buffer_for_geom(
+            geometry_serde_general.GeometryTypeID.POINT,
+            coord_type,
+            8 + geometry_serde_general.CoordinateType.bytes_per_coord(coord_type),
+            1,
+        )
+
+        with pytest.raises(ValueError, match="requires geomserde_speedup"):
+            geometry_serde_general.deserialize(buffer)
+
+    def test_general_serializer_does_not_query_m_on_older_geos(self, monkeypatch):
+        def fail_if_queried(_geometry):
+            raise AssertionError("has_m should not be queried with GEOS < 3.12")
+
+        monkeypatch.setattr(shapely, "geos_version", (3, 11, 0), raising=False)
+        monkeypatch.setattr(
+            BaseGeometry, "has_m", property(fail_if_queried), raising=False
+        )
+        monkeypatch.setattr(
+            geometry_serde_general, "serialize_point", lambda _geometry: b"xy"
+        )
+
+        buffer = geometry_serde_general.serialize(Point(1, 2))
+
+        assert buffer == b"xy"
+
     @staticmethod
     def _test_serde_roundtrip(geoms):
         for geom in geoms: