Update to_parquet to handle custom schema (to fix writing partitions with all missing data) (#201)

Author: jorisvandenbossche

dask_geopandas/io/arrow.py

@@ -191,11 +191,13 @@ def _pandas_to_arrow_table(
 
         table = _geopandas_to_arrow(df, index=preserve_index)
 
-        # TODO add support for schema
-        # (but let it already pass if the passed schema would not change the result)
         if schema is not None:
-            if not table.schema.equals(schema) and len(df):
-                raise NotImplementedError("Passing 'schema' is not yet supported")
+            if not table.schema.equals(schema):
+                # table.schema.metadata contains the "geo" metadata, so
+                # ensure to preserve this in the cast operation
+                if table.schema.metadata and not schema.metadata:
+                    schema = schema.with_metadata(table.schema.metadata)
+                table = table.cast(schema)
 
         return table
 

dask_geopandas/tests/io/test_parquet.py

@@ -161,6 +161,25 @@ def test_parquet_empty_partitions(tmp_path):
     assert result.spatial_partitions is None
 
 
+@pytest.mark.skipif(
+    not Version(dask.__version__) >= Version("2022.06.0"),
+    reason="Only works with dask 2022.06.0 or up",
+)
+def test_parquet_partitions_with_all_missing_strings(tmp_path):
+    df = geopandas.GeoDataFrame(
+        {"col": ["a", "b", None, None]},
+        geometry=geopandas.points_from_xy([0, 1, 2, 3], [0, 1, 2, 3]),
+    )
+    # Creating filtered dask dataframe with at least one empty partition
+    ddf = dask_geopandas.from_geopandas(df, npartitions=2)
+
+    basedir = tmp_path / "dataset"
+    ddf.to_parquet(basedir)
+
+    result = dask_geopandas.read_parquet(basedir)
+    assert_geodataframe_equal(result.compute(), df)
+
+
 @pytest.mark.skipif(
     Version(dask.__version__) < Version("2021.10.0"),
     reason="Only correct error message with dask 2021.10.0 or up",