Spatial partitioning, sorting and shuffling

[benbovy]

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:

• refactor xarray.IndexVariable to allow dask arrays or other lazy arrays More flexible index variables pydata/xarray#8124 (although we can work around this by using regular xarray.Variable objects for now)
• add some sort of generic, out-of-core index in Xarray Low memory/out-of-core index? pydata/xarray#1650, or add our own ad-hoc version here
• compute space-filling curve values from geometries (copied or factored out from dask-geopandas?)
• implement spatial shuffling, probably reusing what is being added in Xarray Add GroupBy.shuffle_to_chunks() pydata/xarray#9320

[dcherian]

spatial shuffling,

Can you add a comment to pydata/xarray#9546 describing the workload and what API might be useful please

[benbovy]

Yes sure, I don't have any precise idea yet but I will do when I'll think more about it.

[martinfleis]

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.

[dcherian]

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.

[benbovy]

To be honest I didn't look closely into either dask_geopandas' GeoDataFrame.spatial_shuffle, dask.array.shuffle or pydata/xarray#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 pydata/xarray#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!

[dcherian]

If the distance values are in-memory then pydata/xarray#9320 should work for you!

ds.groupby_bins(...).shuffle_to_chunks()

[benbovy]

Great! I'll give it a try when I find time.