Handle larger than 32bit distances when creating Hilbert distance and spatially partitioning.

[bernardpazio]

I have come across an issue with spatially partitioning data where the distance between geometry as calculated by hilbert_distance is greater than a 32bit integer. The meta for the returned series is set to 32bit which causes dask to cast down to 32bit when setting index, this causes an overflow and the resulting divisions being out of order. Here is a minimal example to show the issue.

import geopandas as gpd
import dask_geopandas as dgpd
from shapely.geometry import Point

df = gpd.GeoDataFrame({
    'geometry': [Point(-103152.516, -8942.156), Point(118914.500, 1010032.562)]
})

ddf = dgpd.from_geopandas(df, npartitions=2)
ddf.spatial_shuffle()

setting correctly to int64 before passing to dask fixes this issue.

from dask.dataframe.shuffle import set_index
import numpy as np

by = ddf.hilbert_distance()
set_index(ddf, by.astype(np.int64)).compute()

It should be a simple fix to handle these cases, so I've gone ahead and created a branch for your review.

[bernardpazio]

Looks like the failures are not related to my changes.

[TomAugspurger]

This should probably be an option provided by the user, rather than assuming int64. The default should be uint32 for backwards compatability.

[TomAugspurger]

Wouldn't distances already be a uint32? I suspect that the changes will need to be at a lower level (maybe in _encode), and then when we're constructing a Series here the dtype will just be the dtype of whatever ndarray is passed.

[bernardpazio]

Yes you are right, i had missed that this is actually an unsigned int. The reason for the overflow is actually that the value is greater than the max value for int but not uint, and dask simply isn't respecting the dtype.

[bernardpazio]

FYI this problem seems to be only an issue with dask<2022.8 and pandas<1.5

[martinfleis]

FYI this problem seems to be only an issue with dask<2022.8 and pandas<1.5

In this case I suggest we recommend people updating their environment and close this PR.