spatial_shuffle fails when loading from a shapefile

[fbunt]

Calling spatial_shuffle on a GeoDataFrame that originates from a shapefile throws a ValueError. Below is an example that uses the attached shapefile:

>>> import dask_geopandas as dgpd
>>> ddf = dgpd.read_file("example/example.shp", npartitions=1).spatial_shuffle()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/dask_geopandas/expr.py", line 847, in spatial_shuffle
    sorted_ddf = self.set_index(
                 ^^^^^^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/dask_geopandas/expr.py", line 634, in set_index
    ddf = super().set_index(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/dask_expr/_collection.py", line 3463, in set_index
    return new_collection(
           ^^^^^^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/dask_expr/_collection.py", line 4769, in new_collection
    meta = expr._meta
           ^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/functools.py", line 1001, in __get__
    val = self.func(instance)
          ^^^^^^^^^^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/dask_expr/_shuffle.py", line 816, in _meta
    return self.frame._meta.set_index(other, drop=self.drop)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/fred/anaconda3/envs/dgpd-error/lib/python3.11/site-packages/pandas/core/frame.py", line 6173, in set_index
    raise ValueError(
ValueError: Length mismatch: Expected 5 rows, received array of length 0

Environment:

dask                      2024.6.2           pyhd8ed1ab_0    conda-forge
dask-core                 2024.6.2           pyhd8ed1ab_0    conda-forge
dask-expr                 1.1.6              pyhd8ed1ab_0    conda-forge
dask-geopandas            0.4.1              pyhd8ed1ab_0    conda-forge
geopandas                 1.0.0              pyhd8ed1ab_0    conda-forge
geopandas-base            1.0.0              pyha770c72_0    conda-forge
numpy                     2.0.0           py311h1461c94_0    conda-forge
pandas                    2.2.2           py311h14de704_1    conda-forge

example.tar.gz

[JamesParrott]

I've nothing more constructive to contribute I'm afraid, but I have reproduced this error.

[fbunt]

I'm happy to work on a fix but my knowledge of the the dask-expr package is limited. @jorisvandenbossche , do you have any pointers on what could cause a meta mismatch?

[fbunt]

OK, I finally understand why the error is occurring: spatial_shuffle sorts the dataframe by calling the set_index method with a new index generated by the chosen metric. This is delegated up to dask-expr, which tries to determine the resulting meta by basically doing new_meta = self._meta.set_index(new_index._meta). The problem is that the frame's meta has 5 rows of real data (see below) while the new index's meta is empty. This seems obvious now, given the error, but at least I understand what is happening better.

>>> ddf._meta
   OBJECTID_1                                           geometry
0           1  POLYGON ((-2052013.364 1936990.624, -2052017.1...
1           2  POLYGON ((-2036165.45 1974204.268, -2036161.41...
2           3  POLYGON ((-2064485.3 1961858.314, -2064451.517...
3           4  POLYGON ((-2081127.411 1935585.344, -2081136.8...
4           5  POLYGON ((-2058945.274 1948115.633, -2059035.8...

One thing I noticed while looking around is that dask-expr DataFrame and Series objects produce empty meta objects. Given the behavior described above, GeoDataFrame._meta being non-empty seems like a bug since the meta objects will always clash and cause an error.

I went back to an older version of dask-geopandas and GeoDataFrame._meta used to return an empty frame. It seems to me like the solution is to return to producing the old empty frame in _meta. @jorisvandenbossche does this sound reasonable? Also, I'm not familiar enough with the codebase to find where meta is actually being generated. Can anyone point me to it?

[TomAugspurger]

@fbunt I think that https://github.com/geopandas/dask-geopandas/blob/main/dask_geopandas/backends.py#L60-L76 might have what you're looking for.

[fbunt]

Thanks!

[fbunt]

It turns out those functions aren't involved in this issue. The root cause is here. The IO func for shapefiles generates a non-empty meta frame. As far as I can tell, this is the only GeoDataFrame source that produces a non-empty meta object in dask-geopandas. It's older code though so I'm not sure how it hasn't caused issues until now. I'll work on a PR.

[RichardScottOZ]

I just saw this from a geopackage - dask-geopandas 0.4.1 / geopandas-base 1.0.1

[jorisvandenbossche]

@RichardScottOZ geopackage uses the same reader implementation as shapefile, so that should also be fixed by #305

[RichardScottOZ]

Thanks, will try it next week.