Skip to content

Add type literals to engine argument in groupby aggregations #54141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 17, 2023
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 55 additions & 9 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,46 @@ class providing the base-class of operations.
{example}
"""

_groupby_agg_method_engine_template = """
Compute {fname} of group values.

Parameters
----------
numeric_only : bool, default {no}
Include only float, int, boolean columns.

.. versionchanged:: 2.0.0

numeric_only no longer accepts ``None``.

min_count : int, default {mc}
The required number of valid values to perform the operation. If fewer
than ``min_count`` non-NA values are present the result will be NA.

engine : str, default None {e}
* ``'cython'`` : Runs rolling apply through C-extensions from cython.
* ``'numba'`` : Runs rolling apply through JIT compiled code from numba.
Only available when ``raw`` is set to ``True``.
* ``None`` : Defaults to ``'cython'`` or globally setting ``compute.use_numba``

engine_kwargs : dict, default None {ek}
* For ``'cython'`` engine, there are no accepted ``engine_kwargs``
* For ``'numba'`` engine, the engine can accept ``nopython``, ``nogil``
and ``parallel`` dictionary keys. The values must either be ``True`` or
``False``. The default ``engine_kwargs`` for the ``'numba'`` engine is
``{{'nopython': True, 'nogil': False, 'parallel': False}}`` and will be
applied to both the ``func`` and the ``apply`` rolling aggregation.

Returns
-------
Series or DataFrame
Computed {fname} of values within each group.

Examples
--------
{example}
"""

_pipe_template = """
Apply a ``func`` with arguments to this %(klass)s object and return its result.

Expand Down Expand Up @@ -2231,7 +2271,7 @@ def hfunc(bvalues: ArrayLike) -> ArrayLike:
def mean(
self,
numeric_only: bool = False,
engine: str = "cython",
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
):
"""
Expand Down Expand Up @@ -2402,7 +2442,7 @@ def median(self, numeric_only: bool = False):
def std(
self,
ddof: int = 1,
engine: str | None = None,
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
numeric_only: bool = False,
):
Expand Down Expand Up @@ -2511,7 +2551,7 @@ def std(
def var(
self,
ddof: int = 1,
engine: str | None = None,
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
numeric_only: bool = False,
):
Expand Down Expand Up @@ -2909,10 +2949,12 @@ def size(self) -> DataFrame | Series:

@final
@doc(
_groupby_agg_method_template,
_groupby_agg_method_engine_template,
fname="sum",
no=False,
mc=0,
e=None,
ek=None,
example=dedent(
"""\
For SeriesGroupBy:
Expand Down Expand Up @@ -2952,7 +2994,7 @@ def sum(
self,
numeric_only: bool = False,
min_count: int = 0,
engine: str | None = None,
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
):
if maybe_use_numba(engine):
Expand Down Expand Up @@ -3026,10 +3068,12 @@ def prod(self, numeric_only: bool = False, min_count: int = 0):

@final
@doc(
_groupby_agg_method_template,
_groupby_agg_method_engine_template,
fname="min",
no=False,
mc=-1,
e=None,
ek=None,
example=dedent(
"""\
For SeriesGroupBy:
Expand Down Expand Up @@ -3069,7 +3113,7 @@ def min(
self,
numeric_only: bool = False,
min_count: int = -1,
engine: str | None = None,
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
):
if maybe_use_numba(engine):
Expand All @@ -3092,10 +3136,12 @@ def min(

@final
@doc(
_groupby_agg_method_template,
_groupby_agg_method_engine_template,
fname="max",
no=False,
mc=-1,
e=None,
ek=None,
example=dedent(
"""\
For SeriesGroupBy:
Expand Down Expand Up @@ -3135,7 +3181,7 @@ def max(
self,
numeric_only: bool = False,
min_count: int = -1,
engine: str | None = None,
engine: Literal["cython", "numba"] | None = None,
engine_kwargs: dict[str, bool] | None = None,
):
if maybe_use_numba(engine):
Expand Down