Repository: geopandas/dask-geopandas

Issue:
BUG: Metadata inference issue with named geometry column
If you create dask GeoDataFrame from a geopandas gdf with geometry column name different than "geometry", calling `.geometry` in `map_partitions` raises an attribute error. With `"geometry"` name the same code is fine.


```py
import geopandas as gpd
import dask_geopandas

df = gpd.read_file(gpd.datasets.get_path("naturalearth_lowres")).rename_geometry("geom")
ddf = dask_geopandas.from_geopandas(df, npartitions=4)

def get_geometry(partition):
    return partition.geometry

ddf.map_partitions(get_geometry)
```
```py
ValueError: Metadata inference failed in `get_geometry`.

You have supplied a custom function and Dask is unable to 
determine the type of output that that function returns. 

To resolve this please provide a meta= keyword.
The docstring of the Dask function you ran should have more information.

Original error is below:
------------------------
AttributeError("'GeoDataFrame' object has no attribute 'geometry'")

Traceback:
---------
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py", line 176, in raise_on_meta_error
    yield
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5672, in _emulate
    return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
  File "/var/folders/kp/fxnnw89x5qbb9gryn507p03c0000gn/T/ipykernel_52968/1774662833.py", line 2, in get_geometry
    return partition.geometry
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/pandas/core/generic.py", line 5478, in __getattr__
    return object.__getattribute__(self, name)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py in raise_on_meta_error(funcname, udf)
    175     try:
--> 176         yield
    177     except Exception as e:

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _emulate(func, *args, **kwargs)
   5671     with raise_on_meta_error(funcname(func), udf=kwargs.pop("udf", False)):
-> 5672         return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
   5673 

/var/folders/kp/fxnnw89x5qbb9gryn507p03c0000gn/T/ipykernel_52968/1774662833.py in get_geometry(partition)
      1 def get_geometry(partition):
----> 2     return partition.geometry
      3 ddf.map_partitions(get_geometry)

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/pandas/core/generic.py in __getattr__(self, name)
   5477             return self[name]
-> 5478         return object.__getattribute__(self, name)
   5479 

AttributeError: 'GeoDataFrame' object has no attribute 'geometry'

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
/var/folders/kp/fxnnw89x5qbb9gryn507p03c0000gn/T/ipykernel_52968/1774662833.py in <module>
      1 def get_geometry(partition):
      2     return partition.geometry
----> 3 ddf.map_partitions(get_geometry)

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in map_partitions(self, func, *args, **kwargs)
    704         None as the division.
    705         """
--> 706         return map_partitions(func, self, *args, **kwargs)
    707 
    708     @insert_meta_param_description(pad=12)

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in map_partitions(func, meta, enforce_metadata, transform_divisions, *args, **kwargs)
   5726         # Use non-normalized kwargs here, as we want the real values (not
   5727         # delayed values)
-> 5728         meta = _emulate(func, *args, udf=True, **kwargs)
   5729     else:
   5730         meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py in _emulate(func, *args, **kwargs)
   5670     """
   5671     with raise_on_meta_error(funcname(func), udf=kwargs.pop("udf", False)):
-> 5672         return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
   5673 
   5674 

/opt/miniconda3/envs/geo_dev/lib/python3.9/contextlib.py in __exit__(self, type, value, traceback)
    133                 value = type()
    134             try:
--> 135                 self.gen.throw(type, value, traceback)
    136             except StopIteration as exc:
    137                 # Suppress StopIteration *unless* it's the same exception that

/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py in raise_on_meta_error(funcname, udf)
    195         )
    196         msg = msg.format(" in `{0}`".format(funcname) if funcname else "", repr(e), tb)
--> 197         raise ValueError(msg) from e
    198 
    199 

ValueError: Metadata inference failed in `get_geometry`.

You have supplied a custom function and Dask is unable to 
determine the type of output that that function returns. 

To resolve this please provide a meta= keyword.
The docstring of the Dask function you ran should have more information.

Original error is below:
------------------------
AttributeError("'GeoDataFrame' object has no attribute 'geometry'")

Traceback:
---------
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/utils.py", line 176, in raise_on_meta_error
    yield
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/dask/dataframe/core.py", line 5672, in _emulate
    return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))
  File "/var/folders/kp/fxnnw89x5qbb9gryn507p03c0000gn/T/ipykernel_52968/1774662833.py", line 2, in get_geometry
    return partition.geometry
  File "/opt/miniconda3/envs/geo_dev/lib/python3.9/site-packages/pandas/core/generic.py", line 5478, in __getattr__
    return object.__getattribute__(self, name)
```


Discussion:
jorisvandenbossche: This fixes the issue, but still causes a problem in another test (`test_set_geometry_property_on_geodf`):

```
dask_obj.rename(columns={"geometry": "foo"}).set_geometry("foo").compute()
```

If you do this in GeoPandas, after the `rename`, you also have an invalid GeoDataFrame for a moment, and then you "fix" that with the `set_geometry` directly after it. 
But with dask, it's apparently creating a `_meta_nonempty` (I think because `set_geometry` is calling `map_partitions` which runs the function on `_meta_nonempty`, but which at that point thus is an invalid geodataframe that cannot be created ..)

Of course, for this reason it's also better to use `rename_geometry` now. But I suppose the way to fix this is by setting the `meta` manually for `set_geometry`, as we can in general expect that in `set_geometry` you might have an invalid dataframe as you are calling that to set the geometry column.

Patch:
diff --git a/dask_geopandas/backends.py b/dask_geopandas/backends.py
index 6ca0999..89b47c5 100644
--- a/dask_geopandas/backends.py
+++ b/dask_geopandas/backends.py
@@ -57,7 +57,7 @@ def _nonempty_geoseries(x, idx=None):
 @meta_nonempty.register(geopandas.GeoDataFrame)
 def _nonempty_geodataframe(x):
     df = meta_nonempty_dataframe(x)
-    return geopandas.GeoDataFrame(df, crs=x.crs)
+    return geopandas.GeoDataFrame(df, geometry=x._geometry_column_name, crs=x.crs)
 
 
 @make_meta_dispatch.register((geopandas.GeoSeries, geopandas.GeoDataFrame))
diff --git a/dask_geopandas/core.py b/dask_geopandas/core.py
index 69911ee..4f189d2 100644
--- a/dask_geopandas/core.py
+++ b/dask_geopandas/core.py
@@ -11,6 +11,8 @@ from dask.base import tokenize
 import geopandas
 from shapely.geometry.base import BaseGeometry
 from shapely.geometry import box
+import pygeos
+
 from .hilbert_distance import _hilbert_distance
 from .morton_distance import _morton_distance
 
@@ -67,10 +69,6 @@ class _Frame(dd.core._Frame, OperatorMethodMixin):
     def __dask_postpersist__(self):
         return type(self), (self._name, self._meta, self.divisions)
 
-    def __repr__(self):
-        s = "<dask_geopandas.%s | %d tasks | %d npartitions>"
-        return s % (type(self).__name__, len(self.dask), self.npartitions)
-
     @classmethod
     def _bind_property(cls, attr, preserve_spatial_partitions=False):
         """Map property to partitions and bind to class"""
@@ -118,9 +116,14 @@ class _Frame(dd.core._Frame, OperatorMethodMixin):
         # TEMP method to calculate spatial partitions for testing, need to
         # add better methods (set_partitions / repartition)
         parts = geopandas.GeoSeries(
-            self.map_partitions(lambda part: part.convex_hull.unary_union).compute()
+            self.map_partitions(
+                lambda part: pygeos.convex_hull(
+                    pygeos.geometrycollections(part.geometry.values.data)
+                )
+            ).compute(),
+            crs=self.crs,
         )
-        self.spatial_partitions = parts.convex_hull
+        self.spatial_partitions = parts
 
     def _propagate_spatial_partitions(self, new_object):
         """
@@ -413,7 +416,11 @@ class GeoDataFrame(_Frame, dd.core.DataFrame):
         self.dask = new.dask
 
     def set_geometry(self, col):
-        return self.map_partitions(M.set_geometry, col)
+        # calculate ourselves to use meta and not meta_nonempty, which would
+        # raise an error if meta is an invalid GeoDataFrame (e.g. geometry
+        # column name not yet set correctly)
+        meta = self._meta.set_geometry(col)
+        return self.map_partitions(M.set_geometry, col, meta=meta)
 
     def __getitem__(self, key):
         """
diff --git a/dask_geopandas/io/parquet.py b/dask_geopandas/io/parquet.py
index e2ee187..762ad88 100644
--- a/dask_geopandas/io/parquet.py
+++ b/dask_geopandas/io/parquet.py
@@ -58,11 +58,6 @@ class GeoArrowEngine(ArrowEngine):
     def read_metadata(cls, *args, **kwargs):
         meta, stats, parts, index = super().read_metadata(*args, **kwargs)
 
-        # Update meta to be a GeoDataFrame
-        # TODO convert columns based on GEO metadata (this will now only work
-        # for a default "geometry" column)
-        meta = geopandas.GeoDataFrame(meta)
-
         # get spatial partitions if available
         regions = geopandas.GeoSeries([_get_partition_bounds(part) for part in parts])
         if regions.notna().all():
@@ -71,6 +66,23 @@ class GeoArrowEngine(ArrowEngine):
 
         return (meta, stats, parts, index)
 
+    @classmethod
+    def _generate_dd_meta(cls, schema, index, categories, partition_info):
+        meta, index_cols, categories, index, partition_info = super()._generate_dd_meta(
+            schema, index, categories, partition_info
+        )
+        if schema.metadata and b"geo" in schema.metadata:
+            geo_meta = json.loads(schema.metadata[b"geo"])
+            crs = geo_meta["columns"][geo_meta["primary_column"]]["crs"]
+        else:
+            crs = None
+
+        # Update meta to be a GeoDataFrame
+        # TODO convert columns based on GEO metadata (this will now only work
+        # for a default "geometry" column)
+        meta = geopandas.GeoDataFrame(meta, crs=crs)
+        return meta, index_cols, categories, index, partition_info
+
     @classmethod
     def _arrow_table_to_pandas(
         cls, arrow_table: "pyarrow.Table", categories, **kwargs
diff --git a/setup.py b/setup.py
index d53add8..48106cd 100644
--- a/setup.py
+++ b/setup.py
@@ -8,6 +8,7 @@ install_requires = [
     "dask>=2.18.0,!=2021.05.1",
     "distributed>=2.18.0,!=2021.05.1",
     "numba",
+    "pygeos",
 ]
 
 


Test patch:
diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py
index 738fbbe..4e0d34c 100644
--- a/tests/io/test_parquet.py
+++ b/tests/io/test_parquet.py
@@ -30,6 +30,8 @@ def test_parquet_roundtrip(tmp_path):
     result = dask_geopandas.read_parquet(basedir)
     assert result.npartitions == 4
     assert_geodataframe_equal(result.compute(), df)
+    # reading back correctly sets the CRS in meta
+    assert result.crs == df.crs
     # reading back also populates the spatial partitioning property
     assert result.spatial_partitions is not None
 
diff --git a/tests/test_core.py b/tests/test_core.py
index 2ff7a94..e36606d 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -334,6 +334,30 @@ def test_set_geometry_property_on_geodf(geodf_points):
     assert all(df.geometry == df.foo)
 
 
+def test_meta(geodf_points_crs):
+    df = geodf_points_crs
+    dask_obj = dask_geopandas.from_geopandas(df, npartitions=2)
+
+    def check_meta(gdf, name):
+        assert isinstance(gdf, geopandas.GeoDataFrame)
+        assert isinstance(gdf.geometry, geopandas.GeoSeries)
+        assert gdf.crs == df.crs
+        assert gdf._geometry_column_name == name
+
+    meta = dask_obj._meta
+    check_meta(meta, "geometry")
+    meta_non_empty = dask_obj._meta_nonempty
+    check_meta(meta_non_empty, "geometry")
+
+    # with non-default geometry name
+    df = df.rename_geometry("foo")
+    dask_obj = dask_geopandas.from_geopandas(df, npartitions=2)
+    meta = dask_obj._meta
+    check_meta(meta, "foo")
+    meta_non_empty = dask_obj._meta_nonempty
+    check_meta(meta_non_empty, "foo")
+
+
 def test_to_crs_geodf(geodf_points_crs):
     df = geodf_points_crs
     dask_obj = dask_geopandas.from_geopandas(df, npartitions=2)
@@ -394,6 +418,21 @@ def test_geoseries_apply(geoseries_polygons):
     pd.testing.assert_series_equal(result, expected)
 
 
-def test_geodataframe_html_repr(geodf_points):
+def test_repr(geodf_points):
     dask_obj = dask_geopandas.from_geopandas(geodf_points, npartitions=2)
+    assert "Dask GeoDataFrame" in dask_obj.__repr__()
+    assert "Dask GeoSeries" in dask_obj.geometry.__repr__()
     assert "Dask-GeoPandas GeoDataFrame" in dask_obj._repr_html_()
+
+
+def test_map_partitions_get_geometry(geodf_points):
+    # https://github.com/geopandas/dask-geopandas/issues/100
+    df = geodf_points.rename_geometry("foo")
+    dask_obj = dask_geopandas.from_geopandas(df, npartitions=2)
+
+    def get_geometry(partition):
+        return partition.geometry
+
+    result = dask_obj.map_partitions(get_geometry).compute()
+    expected = dask_obj.geometry.compute()
+    assert_geoseries_equal(result, expected)