Don't put GeoSeries in map_partitions kwarg (#205)

* Workaround for the fact that we can't include a partitioned _Frame in
map_partitions kwargs without it being concatenated on a single worker.
This is an upstream bug dask/dask#8308, but
the fix is challenging.

* Check for dd.core._Frame

* Add test

* Address code review

Author: ian-r-rose

dask_geopandas/core.py

@@ -778,6 +778,15 @@ def from_dask_dataframe(df, geometry=None):
         the values will be set as 'geometry' column on the GeoDataFrame.
 
     """
+    # If the geometry column is already a partitioned `_Frame`, we can't refer to
+    # it via a keyword-argument due to https://github.com/dask/dask/issues/8308.
+    # Instead, we assign the geometry column using regular dataframe operations,
+    # then refer to that column by name in `map_partitions`.
+    if isinstance(geometry, dd.core.Series):
+        name = geometry.name if geometry.name is not None else "geometry"
+        return df.assign(**{name: geometry}).map_partitions(
+            geopandas.GeoDataFrame, geometry=name
+        )
     return df.map_partitions(geopandas.GeoDataFrame, geometry=geometry)
 
 

dask_geopandas/tests/test_core.py

@@ -390,6 +390,12 @@ def test_from_dask_dataframe_with_dask_geoseries():
     dask_obj = dask_geopandas.from_dask_dataframe(
         dask_obj, geometry=dask_geopandas.points_from_xy(dask_obj, "x", "y")
     )
+    # Check that the geometry isn't concatenated and embedded a second time in
+    # the high-level graph. cf. https://github.com/geopandas/dask-geopandas/issues/197
+    k = next(k for k in dask_obj.dask.dependencies if k.startswith("GeoDataFrame"))
+    deps = dask_obj.dask.dependencies[k]
+    assert len(deps) == 1
+
     expected = df.set_geometry(geopandas.points_from_xy(df["x"], df["y"]))
     assert_geoseries_equal(dask_obj.geometry.compute(), expected.geometry)