Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Spatial partitioning, sorting and shuffling #84

Open
benbovy opened this issue Sep 27, 2024 · 7 comments
Open

Spatial partitioning, sorting and shuffling #84

benbovy opened this issue Sep 27, 2024 · 7 comments

Comments

@benbovy
Copy link
Member

benbovy commented Sep 27, 2024

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:

@dcherian
Copy link
Contributor

spatial shuffling,

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

@benbovy
Copy link
Member Author

benbovy commented Sep 27, 2024

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

@martinfleis
Copy link
Member

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
Copy link
Contributor

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
Copy link
Member Author

benbovy commented Sep 30, 2024

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
Copy link
Contributor

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

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

@benbovy
Copy link
Member Author

benbovy commented Nov 22, 2024

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants