PERF: Allow groupby transform with numba engine to be fully parallelizable

[mroeschke]

• tests added / passed
• passes black pandas
• passes git diff upstream/master -u -- "*.py" | flake8 --diff
• whatsnew entry

New performance comparison with 10k groups

In [1]: df_g = pd.DataFrame({'a': range(10**4), 'b': range(10**4), 'c': range(10**4)})

In [2]: df_g.groupby('a').transform(lambda x: x + 1)
Out[2]:
          b      c
0         1      1
1         2      2
2         3      3
3         4      4
4         5      5
...     ...    ...
9995   9996   9996
9996   9997   9997
9997   9998   9998
9998   9999   9999
9999  10000  10000

[10000 rows x 2 columns]

In [3]: %timeit df_g.groupby('a').transform(lambda x: x + 1)
9.48 s ± 885 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [4]: def f(values, index):
   ...:     return values + 1
   ...:

In [5]: df_g.groupby('a').transform(f, engine='numba', engine_kwargs={'parallel': True})
Out[5]:
            b        c
0         1.0      1.0
1         2.0      2.0
2         3.0      3.0
3         4.0      4.0
4         5.0      5.0
...       ...      ...
9995   9996.0   9996.0
9996   9997.0   9997.0
9997   9998.0   9998.0
9998   9999.0   9999.0
9999  10000.0  10000.0

[10000 rows x 2 columns]

In [6]: %timeit df_g.groupby('a').transform(f, engine='numba', engine_kwargs={'parallel': True})
4.41 ms ± 34.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[jbrockmendel]

How does this compare to non-parallel numba?

[mroeschke]

How does this compare to non-parallel numba?

In [1]: df_g = pd.DataFrame({'a': range(10**4), 'b': range(10**4), 'c': range(10**4)})

In [2]: In [4]: def f(values, index):
   ...:    ...:     return values + 1
   ...:

In [3]: df_g.groupby('a').transform(f, engine='numba')
Out[3]:
            b        c
0         1.0      1.0
1         2.0      2.0
2         3.0      3.0
3         4.0      4.0
4         5.0      5.0
...       ...      ...
9995   9996.0   9996.0
9996   9997.0   9997.0
9997   9998.0   9998.0
9998   9999.0   9999.0
9999  10000.0  10000.0

[10000 rows x 2 columns]

In [4]: %timeit df_g.groupby('a').transform(f, engine='numba')
5.46 ms ± 38.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[jreback]

so a little better :)

[jreback]

can you add an asv which covers this (or do we have an existing one)?

[jreback]

maybe this check should actually be in _transform_with_numba to keep DRY (you have it above as well)

[jreback]

thanks @mroeschke