ENH: add automatic detection of 2.5D in write_dataframe (#223)

Author: theroggy

CHANGES.md

@@ -4,7 +4,8 @@
 
 ### Improvements
 
-- ENH: add "driver" property to read_info result (#224)
+- Add automatic detection of 2.5D geometries in write_dataframe (#223)
+- Add "driver" property to read_info result (#224)
 
 ## 0.5.1 (2023-01-26)
 

pyogrio/geopandas.py

@@ -1,5 +1,7 @@
-from pyogrio.raw import DRIVERS_NO_MIXED_SINGLE_MULTI
+import numpy as np
+from pyogrio.raw import DRIVERS_NO_MIXED_SINGLE_MULTI, DRIVERS_NO_MIXED_DIMENSIONS
 from pyogrio.raw import detect_driver, read, read_arrow, write
+from pyogrio.errors import DataSourceError
 
 
 def _stringify_path(path):
@@ -309,9 +311,22 @@ def write_dataframe(
     # Determine geometry_type and/or promote_to_multi
     if geometry_type is None or promote_to_multi is None:
         tmp_geometry_type = "Unknown"
+        has_z = False
 
         # If there is data, infer layer geometry type + promote_to_multi
         if not df.empty:
+            # None/Empty geometries sometimes report as Z incorrectly, so ignore them
+            has_z_arr = geometry[
+                (geometry != np.array(None)) & (~geometry.is_empty)
+            ].has_z
+            has_z = has_z_arr.any()
+            all_z = has_z_arr.all()
+
+            if driver in DRIVERS_NO_MIXED_DIMENSIONS and has_z and not all_z:
+                raise DataSourceError(
+                    f"Mixed 2D and 3D coordinates are not supported by {driver}"
+                )
+
             geometry_types = pd.Series(geometry.type.unique()).dropna().values
             if len(geometry_types) == 1:
                 tmp_geometry_type = geometry_types[0]
@@ -347,6 +362,8 @@ def write_dataframe(
 
         if geometry_type is None:
             geometry_type = tmp_geometry_type
+            if has_z:
+                geometry_type = f"2.5D {geometry_type}"
 
     crs = None
     if geometry.crs:

pyogrio/raw.py

@@ -32,6 +32,10 @@
     "GPKG",
 }
 
+DRIVERS_NO_MIXED_DIMENSIONS = {
+    "FlatGeobuf",
+}
+
 
 def read(
     path_or_buffer,

pyogrio/tests/test_geopandas_io.py

@@ -7,14 +7,19 @@
 from pyogrio import list_layers, read_info, __gdal_version__, __gdal_geos_version__
 from pyogrio.errors import DataLayerError, DataSourceError, FeatureError, GeometryError
 from pyogrio.geopandas import read_dataframe, write_dataframe
-from pyogrio.raw import DRIVERS, DRIVERS_NO_MIXED_SINGLE_MULTI
+from pyogrio.raw import (
+    DRIVERS,
+    DRIVERS_NO_MIXED_DIMENSIONS,
+    DRIVERS_NO_MIXED_SINGLE_MULTI,
+)
 from pyogrio.tests.conftest import ALL_EXTS
 
 try:
     import pandas as pd
     from pandas.testing import assert_frame_equal, assert_index_equal
 
     import geopandas as gp
+    from geopandas.array import from_wkt
     from geopandas.testing import assert_geodataframe_equal
 
     from shapely.geometry import Point
@@ -843,17 +848,93 @@ def test_write_read_null(tmp_path):
     ],
 )
 def test_write_geometry_z_types(tmp_path, wkt, geom_types):
-    from geopandas.array import from_wkt
-
     filename = tmp_path / "test.fgb"
-
     gdf = gp.GeoDataFrame(geometry=from_wkt([wkt]), crs="EPSG:4326")
     for geom_type in geom_types:
         write_dataframe(gdf, filename, geometry_type=geom_type)
         df = read_dataframe(filename)
         assert_geodataframe_equal(df, gdf)
 
 
+@pytest.mark.parametrize("ext", ALL_EXTS)
+@pytest.mark.parametrize(
+    "test_descr, exp_geometry_type, mixed_dimensions, wkt",
+    [
+        ("1 Point Z", "2.5D Point", False, ["Point Z (0 0 0)"]),
+        ("1 LineString Z", "2.5D LineString", False, ["LineString Z (0 0 0, 1 1 0)"]),
+        (
+            "1 Polygon Z",
+            "2.5D Polygon",
+            False,
+            ["Polygon Z ((0 0 0, 0 1 0, 1 1 0, 0 0 0))"],
+        ),
+        ("1 MultiPoint Z", "2.5D MultiPoint", False, ["MultiPoint Z (0 0 0, 1 1 0)"]),
+        (
+            "1 MultiLineString Z",
+            "2.5D MultiLineString",
+            False,
+            ["MultiLineString Z ((0 0 0, 1 1 0), (2 2 2, 3 3 2))"],
+        ),
+        (
+            "1 MultiLinePolygon Z",
+            "2.5D MultiPolygon",
+            False,
+            [
+                "MultiPolygon Z (((0 0 0, 0 1 0, 1 1 0, 0 0 0)), ((1 1 1, 1 2 1, 2 2 1, 1 1 1)))"  # noqa: E501
+            ],
+        ),
+        (
+            "1 GeometryCollection Z",
+            "2.5D GeometryCollection",
+            False,
+            ["GeometryCollection Z (Point Z (0 0 0))"],
+        ),
+        ("Point Z + Point", "2.5D Point", True, ["Point Z (0 0 0)", "Point (0 0)"]),
+        ("Point Z + None", "2.5D Point", False, ["Point Z (0 0 0)", None]),
+    ],
+)
+def test_write_geometry_z_types_auto(
+    tmp_path, ext, test_descr, exp_geometry_type, mixed_dimensions, wkt
+):
+    # Shapefile has some different behaviour that other file types
+    if ext == ".shp":
+        if exp_geometry_type == "2.5D GeometryCollection":
+            pytest.skip(f"ext {ext} doesn't support {exp_geometry_type}")
+        elif exp_geometry_type == "2.5D MultiLineString":
+            exp_geometry_type = "2.5D LineString"
+        elif exp_geometry_type == "2.5D MultiPolygon":
+            exp_geometry_type = "2.5D Polygon"
+
+    column_data = {}
+    column_data["test_descr"] = [test_descr] * len(wkt)
+    column_data["idx"] = [str(idx) for idx in range(len(wkt))]
+    gdf = gp.GeoDataFrame(column_data, geometry=from_wkt(wkt), crs="EPSG:4326")
+    filename = tmp_path / f"test{ext}"
+
+    if mixed_dimensions and DRIVERS[ext] in DRIVERS_NO_MIXED_DIMENSIONS:
+        with pytest.raises(
+            DataSourceError,
+            match=("Mixed 2D and 3D coordinates are not supported by"),
+        ):
+            write_dataframe(gdf, filename)
+        return
+    else:
+        write_dataframe(gdf, filename)
+
+    info = read_info(filename)
+    assert info["geometry_type"] == exp_geometry_type
+
+    result_gdf = read_dataframe(filename)
+    if ext == ".geojsonl":
+        result_gdf.crs = "EPSG:4326"
+    if ext == ".fgb":
+        # When the following gdal issue is released, this if needs to be removed:
+        # https://github.com/OSGeo/gdal/issues/7401
+        gdf = gdf.loc[~((gdf.geometry == np.array(None)) | gdf.geometry.is_empty)]
+
+    assert_geodataframe_equal(gdf, result_gdf)
+
+
 def test_read_multisurface(data_dir):
     df = read_dataframe(data_dir / "test_multisurface.gpkg")