BUG: fix issue reading with use_arrow=True after having read a Parquet file (#601)

Author: jorisvandenbossche

CHANGES.md

@@ -18,6 +18,8 @@
 -   Fix decode error reading an sqlite file on windows (#568).
 -   Fix wrong layername when creating .gpkg.zip file (#570).
 -   Fix segfault on providing an invalid value for `layer` in `read_info` (#564).
+-   Fix error when reading data with ``use_arrow=True`` after having used the
+    Parquet driver with GDAL>=3.12 (#601).
 
 ### Packaging
 

ci/envs/latest.yml

@@ -4,6 +4,7 @@ channels:
 dependencies:
   - numpy
   - libgdal-core
+  - libgdal-arrow-parquet
   - pytest
   - shapely>=2
   - geopandas-base

pyogrio/geopandas.py

@@ -8,6 +8,7 @@
 
 from pyogrio._compat import (
     HAS_GEOPANDAS,
+    HAS_PYARROW,
     PANDAS_GE_15,
     PANDAS_GE_20,
     PANDAS_GE_22,
@@ -346,8 +347,18 @@ def read_dataframe(
         elif geometry_name in df.columns:
             wkb_values = df.pop(geometry_name)
             if PANDAS_GE_15 and wkb_values.dtype != object:
-                # for example ArrowDtype will otherwise create numpy array with pd.NA
-                wkb_values = wkb_values.to_numpy(na_value=None)
+                if (
+                    HAS_PYARROW
+                    and isinstance(wkb_values.dtype, pd.ArrowDtype)
+                    and isinstance(wkb_values.dtype.pyarrow_dtype, pa.BaseExtensionType)
+                ):
+                    # handle BaseExtensionType(extension<geoarrow.wkb>)
+                    wkb_values = pa.array(wkb_values.array).to_numpy(
+                        zero_copy_only=False
+                    )
+                else:
+                    # for example ArrowDtype will otherwise give numpy array with pd.NA
+                    wkb_values = wkb_values.to_numpy(na_value=None)
             df["geometry"] = shapely.from_wkb(wkb_values, on_invalid=on_invalid)
             if force_2d:
                 df["geometry"] = shapely.force_2d(df["geometry"])

pyogrio/tests/conftest.py

@@ -51,6 +51,8 @@
     ".shp": 0,
 }
 
+GDAL_HAS_PARQUET_DRIVER = "Parquet" in list_drivers()
+
 
 def pytest_report_header(config):
     drivers = ", ".join(

pyogrio/tests/test_geopandas_io.py

@@ -35,6 +35,7 @@
 from pyogrio.tests.conftest import (
     ALL_EXTS,
     DRIVERS,
+    GDAL_HAS_PARQUET_DRIVER,
     START_FID,
     requires_arrow_write_api,
     requires_gdal_geos,
@@ -2644,3 +2645,21 @@ def test_write_geojson_rfc7946_coordinates(tmp_path, use_arrow):
 
     gdf_in_appended = read_dataframe(output_path, use_arrow=use_arrow)
     assert np.array_equal(gdf_in_appended.geometry.values, points + points_append)
+
+
+@pytest.mark.requires_arrow_api
+@pytest.mark.skipif(
+    not GDAL_HAS_PARQUET_DRIVER, reason="Parquet driver is not available"
+)
+def test_parquet_driver(tmp_path, use_arrow):
+    """
+    Simple test verifying the Parquet driver works if available
+    """
+    gdf = gp.GeoDataFrame(
+        {"col": [1, 2, 3], "geometry": [Point(0, 0), Point(1, 1), Point(2, 2)]},
+        crs="EPSG:4326",
+    )
+    output_path = tmp_path / "test.parquet"
+    write_dataframe(gdf, output_path, use_arrow=use_arrow)
+    result = read_dataframe(output_path, use_arrow=use_arrow)
+    assert_geodataframe_equal(result, gdf)