[Data] Optimizing the multi column groupby

[wingkitlee0]

This PR improves the "get group boundaries" function in groupby.map_groups(). This is not a bottleneck so the overall runtime will likely have minimal improvement, but the function itself is optimized and the code is cleaner.

Changes:

• Instead of custom class implementation (_MultiColumnSortedKey, which I wrote a few months ago...), a numpy record array is constructed.
• Changed from "np.searchsorted + while-loop" to np.where

In addition, I have cleaned some indexing and added docstrings and tests.

Why are these changes needed?

The pure numpy implementation is about 8x faster. The result below is for 100K elements. (similar results for 10M elements)

------------------------------------------------------------------------------------------------ benchmark: 2 tests -----------------------------------------------------------------------------------------------
Name (time in ms)                                    Min                 Max                Mean            StdDev              Median               IQR            Outliers      OPS            Rounds  Iterations
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_benchmark[get_key_boundaries]               21.1100 (1.0)       25.5830 (1.0)       22.2432 (1.0)      1.0147 (1.0)       22.1691 (1.0)      0.8943 (1.0)           8;4  44.9575 (1.0)          30           1
test_benchmark[get_key_boundaries_original]     175.2830 (8.30)     178.6610 (6.98)     176.4963 (7.93)     1.4369 (1.42)     175.8079 (7.93)     2.5007 (2.80)          2;0   5.6658 (0.13)          6           1
-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

The custom _MultiColumnSortedKey class was introduced as a work-around. The numpy recarray should be preferred as it allows vectorization.

Time complexity / Speed up

Most of the speed up is coming from the use of recarray. This gives about 7x speed up.

A little more speed up was achieved when using np.where, although the time complexity changed from the current O( k log n) to O(n) with the new implementation (np.where), where k is the number of groups and n is the length of the array. This is likely due to the vectorization (less branching).

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[wingkitlee0]

Hey @bveeramani , just wanna put this PR under your radar :)

[bveeramani]

Hey @wingkitlee0, thanks for the ping! I'll let our oncall know (he's the one who's responsible for reviewing external PRs)

[wingkitlee0]

not sure how to solve Read the Docs build failed!. Any idea?

[wingkitlee0]

Hi @alexeykudinkin this PR is ready

[wingkitlee0]

@alexeykudinkin the PR is ready for review again; added nan/None tests

[wingkitlee0]

bump

[alexeykudinkin]

LGTM, minor comments