Accelerate coordinate selection via keysort

[alimanfoo]

There is a potential optimisation for get/set_coordinate_selection() by making use of the keysort() function implemented in pytables, see this comment. Pytables is too heavy a dependency but it would be reasonable to factor out the keysort() function into a standalone package then depend on that (e.g., zarr-developers/numpy-keysort or something like that). Because keysort() requires cython to build I'd rather not include within zarr as zarr is currently pure Python and that makes development easier.

[jakirkham]

cc @FrancescAlted

[FrancescAlted]

Another possibility would be to convince the NumPy crew that keysort() would be a nice feature to implement on it. cc @charris, @jaimefrio

[jakirkham]

First step for convincing NumPy would be dropping it on their mailing list and seeing what reception it gets.

On a different note, have found searchsorted very effective in various indexing applications. This is probably not new information for you though and may not apply to the use cases you have in mind. Though if there are tweaks that can be made to searchsorted that would help, that might be a faster path to NumPy acceptance.

[alimanfoo]

Thanks @jakirkham. I use searchsorted a lot too but IIRC it's not applicable in this situation (although don't hold me to that as the indexing code is reasonably involved and I haven't looked at in a while). Specifically you want to do an argsort and also sort at the same time. FWIW this is not critical for me at the moment, performance in my use cases is fine without this additional optimisation, but happy for someone to pursue this if it becomes important for them.

[jakirkham]

@hameerabbasi, saw you were playing with different sorts like radix sort ( particularly in issue pydata/sparse#62 ). Would you be interested in collaborating on this problem? Maybe we can get a nice small sorting library that handles our needs. :)

[jakirkham]

This SO question is pretty interesting.

[hameerabbasi]

Re: Key-sort, it's very easy to do a key sort with existing NumPy functions, and AFAICT you won't get a large speed boost by writing another implementation, especially if a large number of columns are to be sorted by the same index. It's a simple argsort plus indexing with the result of that argsort. That said, NumPy already implements argsort by doing a key-sort on numpy.arange(...).

As for radix sort, I agree it would be nice if it was in NumPy, I needed it myself as @jakirkham pointed out. There was an issue about it (numpy/numpy#6050) but it didn't come to any fruition. I started writing an implementation somewhere, but it seems all of NumPy's sorts are in-place and there would need to be a special case in the function for the temporary array that radix sort needs.

That said, I agree that these things are very broadly useful and I think we should work with NumPy to get them in, and I am quite sure if we made a PR it wouldn't be turned down. There's a very nice and simple implementation of radix sort (which can be used with or without keys) at eloj/radix-sorting.

[alimanfoo]

Thanks all for looking at this. FWIW using an optimised implementation of keysort (currently in pytables) would accelerate coordinate selection in zarr by a factor of ~2 over keysort with numpy via argsort, see link in issue description. So optimised keysort implementation would be valuable IMO either in numpy or a separate package.

[hameerabbasi]

It's an integer sort, so I suspect PyTables is using radix sort internally, which would explain the performance difference. NumPy uses introsort.

[hameerabbasi]

I'm in the process of adding radix sort to NumPy. You can see the PR: numpy/numpy#12586

[jakirkham]

Given that NumPy's sorting was significantly improved in 1.17.0, it is probably worth reassessing whether this is still an issue.

[jakirkham]

Going to go ahead and close as I believe NumPy 1.17.0 should have fixed this, but please feel free to reopen if I've missed something.