xarray-contrib / flox

Fast & furious GroupBy operations for dask.array
https://flox.readthedocs.io
Apache License 2.0
124 stars 18 forks source link

possible support for sparse arrays #295

Open dcherian opened 11 months ago

dcherian commented 11 months ago

@ilan-gold @ivirshup if you have time, it'd be nice to see a groupby-reduce workflow you'd like to see supported natively by flox.

ivirshup commented 11 months ago

Thanks for opening the issue!

I think the workflows we'd like to see are pretty straight forward. We'd like:

I've thought a little bit about implementation.

seberg commented 11 months ago

I recall @seberg working on some ufunc methods for sparse arrays at a sprint, but don't recall if it ended up working out

That mostly worked, but was a bit slow. The point being that the approach was to extract the result sparsity pattern. Then extract the data and apply the normal ufunc to it. But at least for memory bound ufuncs like add, it was much slower (maybe 3-4x, but don't recall). Now, of course that could probably be optimized a bit by specialization e.g. of binary ufuncs. But in the end, the way scipy-sparse works, it seemed potentially useful as a fallback to implement any ufunc for any dtype NumPy supports, even if scipy sparse has never heard of them. But probably not as a replacement for most operations.