Open benbovy opened 3 days ago
spatial shuffling,
Can you add a comment to https://github.com/pydata/xarray/issues/9546 describing the workload and what API might be useful please
Yes sure, I don't have any precise idea yet but I will do when I'll think more about it.
Doesn't this just mean computing Hilbert distance (we can use vanilla Geopandas if needed or vendor that code) and using sortby
along that? That way you can partition the data along the spatial order without any changes on xarray side.
That is also how dask_geopandas
does that and how sort_values(by="geometry")
works in geopandas. But I may as well be missing something.
sort_values(by="geometry")
yeah probably best to reuse dask-geopandas here, we'd need some new dask array API to shuffle by unknown values.
To be honest I didn't look closely into either dask_geopandas' GeoDataFrame.spatial_shuffle
, dask.array.shuffle
or https://github.com/pydata/xarray/issues/9546.
I have looked into that a bit more now but probably not enough yet.
we'd need some new dask array API to shuffle by unknown values.
@dcherian -- What do you mean exactly by "unknown values"?
What we want to shuffle here are "distance" values (encoded as integer indices) along a 1-dimensional space filling curve of a fixed resolution (level
), after computing it from the geometries. For very coarse resolutions we could probably use those integer indices directly as group labels, but since the range of possible index values grows very rapidly with level
we might better want binning the data (where I guess the number of bins and their extents could be calculated from the chunks).
Would it be possible to use something like ds.shuffle(ds.groupby_bins(...))
with https://github.com/pydata/xarray/pull/9320?
Doesn't this just mean computing Hilbert distance (we can use vanilla Geopandas if needed or vendor that code) and using sortby along that? That way you can partition the data along the spatial order without any changes on xarray side. That is also how dask_geopandas does that and how sort_values(by="geometry") works in geopandas. But I may as well be missing something.
@martinfleis -- Hmm from what I see dask_geopandas.GeoDataFrame.spatial_shuffle
relies on dask.dataframe.DataFrame.set_index
, which involves complex sorting / shuffling logic implemented in https://github.com/dask/dask/blob/main/dask/dataframe/shuffle.py.
I guess we might as well reuse dask.dataframe
here (i.e., convert dask.array.Array
<-> dask.dataframe.DataFrame
) or reuse dask-geopandas.
Computing Hilbert distance is more embarrassingly parallel, if we can do it with vanilla Geopandas we should do it!
When dealing with large sets of geometries it would be nice if we could partition (chunk) the geometry coordinate and the GeometryIndex based on spatial locality (thus requiring spatial sorting or shuffling), like explained for geo-dataframes in dask-geopandas' spatial partitioning user guide.
This would require a good amount of work both here and upstream, though:
xarray.IndexVariable
to allow dask arrays or other lazy arrays https://github.com/pydata/xarray/pull/8124 (although we can work around this by using regularxarray.Variable
objects for now)