ENH: Add support for engine=numba to groupby apply
State: closed

- [ ] closes #53933  (Replace xxxx with the GitHub issue number)
- [ ] [Tests added and passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#writing-tests) if fixing a bug or adding a new feature
- [ ] All [code checks passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit).
- [ ] Added [type annotations](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#type-hints) to new arguments/methods/functions.
- [ ] Added an entry in the latest `doc/source/whatsnew/vX.X.X.rst` file if fixing a bug or adding a new feature.

I ended up just wrapping the DF/Series objects so they're recognized by numba.

Currently, all you can do with the DF/Series objects is extract the values and the index from them.

I think it would make sense to implement things like the addition operators, and some of the common reductions on them, but that's for the future.

I went this way since going through the values, index way of the other functions was just way to much of a hassle, and would cause a lot of code duplication and re-implementation. More details in #53933 
This way also lets users write a single function for the "python"/"numba" engines.

Currently done:
- Reductions (to scalar, axis 0, axis 1)
- Transforms ( tested with normalization func, and identity func)
- Expansions (the example in the user guide is tested)

TODO:
- parallel support
   - numba containers are not thread safe. Will look into threading manually with Python and hope the GIL doesn't get in the way (for a follow up).
- Series support in groupby

Labels:
name: Groupby
name: numba
description: numba-accelerated operations
name: Stale

Comments:
OK, this should be ready for review now.

As of right now, this PR does 2 things

1. Allows using pandas Index/Series/DataFrame within numba jitted code.
   a. I can split this part into another PR which should reduce the diff down to 100-200 lines.
2. Allows passing a function to jit with numba into groupby.apply

The API for this is currently really clunky since as of right now, a lot of the pandas API is not implemented yet.
(Basically only the constructors are implemented).

(You can see how you would write a function for the numba engine in ``pandas/tests/numba/test_basic.py``, and ``pandas/tests/groupby/test_numba.py``)

However, all the basic operations I can think of (reductions, transformations, adding a column) should work right now, so it's at least usable in the current state.

And, looking on the bright side, it can only get better from here :) .

Right now, I'm not really looking for a deep review, I'm just trying to get confirmation that this is the right approach.
(I'll clean everything up after that).
Looked through it and made some superficial comments, don't really have much to add.  I'm not a huge fan of the engine/engine_kwargs pattern (xref #53239) so am not enthusiastic about this, but not opposed.  

Does this make any implicit assumptions about having all-numpy dtypes?

Would there be a viable path for this to live in numba or numba-foo (and not be our problem)?
Thanks for the review.

> Looked through it and made some superficial comments, don't really have much to add. I'm not a huge fan of the engine/engine_kwargs pattern (xref #53239) so am not enthusiastic about this, but not opposed.
> 
> Does this make any implicit assumptions about having all-numpy dtypes?

This assumes all numpy for now. I don't think it's too hard to implement this for Arrow, you just need to know the Arrow columnar format well.

(e.g. see https://uwekorn.com/2018/08/03/use-numba-to-work-with-apache-arrow-in-pure-python.html)

> Would there be a viable path for this to live in numba or numba-foo (and not be our problem)?

The way extensions are registered for numba(they use the entrypoints mechanism like we do for plotting), we can definitely move ``pandas/core/_numba/extensions.py`` outside to a separate repo.

Two things we'd need to figure out if we go down this road are:

1) Using pandas internals (Cython code).
    Numba provides a mechanism to call Cython functions from Python code, which could be useful for implementing indexing.

2) Relationship between pandas and this library

     We'd also need to take a optional dependency on the split off library or provide some sort of plugin structure for things 
     such as ``Groupby.apply``, since the looping part is not the same for the Python/numba engines.

I'll add something to the agenda for the dev meeting on Wednesday.
> We'd also need to take a optional dependency on the split off library or provide some sort of plugin structure for things
such as Groupby.apply, since the looping part is not the same for the Python/numba engines.

Unless we expect more of these to pop up, then a plugin structure sounds less appealing than having this live inside pandas.  So pls ignore the previous question about having this live elsewhere.
This pull request is stale because it has been open for thirty days with no activity. Please [update](https://pandas.pydata.org/pandas-docs/stable/development/contributing.html#updating-your-pull-request) and respond to this comment if you're still interested in working on this.
Sorry for the lack of update.

I talked to Graham Markall from Nvidia about this and I think cuDF is using a numpy record dtype to represent a dtype.
(The advantage there is that it supports indexing to get a column and also that it allows for heterogenous dtypes)

I'm currently blocked on a bug in numba I think where an array of records isn't getting unboxed correctly.

For now, I think it's OK to mothball this and focus on DataFrame.apply.