-
-
Notifications
You must be signed in to change notification settings - Fork 17.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
PERF: Improve performance in rolling.mean(engine="numba") (#43612)
- Loading branch information
Showing
7 changed files
with
250 additions
and
21 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
from __future__ import annotations | ||
|
||
from typing import Callable | ||
|
||
import numpy as np | ||
|
||
from pandas._typing import Scalar | ||
from pandas.compat._optional import import_optional_dependency | ||
|
||
from pandas.core.util.numba_ import ( | ||
NUMBA_FUNC_CACHE, | ||
get_jit_arguments, | ||
) | ||
|
||
|
||
def generate_shared_aggregator( | ||
func: Callable[..., Scalar], | ||
engine_kwargs: dict[str, bool] | None, | ||
cache_key_str: str, | ||
): | ||
""" | ||
Generate a Numba function that loops over the columns 2D object and applies | ||
a 1D numba kernel over each column. | ||
Parameters | ||
---------- | ||
func : function | ||
aggregation function to be applied to each column | ||
engine_kwargs : dict | ||
dictionary of arguments to be passed into numba.jit | ||
cache_key_str: str | ||
string to access the compiled function of the form | ||
<caller_type>_<aggregation_type> e.g. rolling_mean, groupby_mean | ||
Returns | ||
------- | ||
Numba function | ||
""" | ||
nopython, nogil, parallel = get_jit_arguments(engine_kwargs, None) | ||
|
||
cache_key = (func, cache_key_str) | ||
if cache_key in NUMBA_FUNC_CACHE: | ||
return NUMBA_FUNC_CACHE[cache_key] | ||
|
||
numba = import_optional_dependency("numba") | ||
|
||
@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) | ||
def column_looper( | ||
values: np.ndarray, | ||
start: np.ndarray, | ||
end: np.ndarray, | ||
min_periods: int, | ||
): | ||
result = np.empty((len(start), values.shape[1]), dtype=np.float64) | ||
for i in numba.prange(values.shape[1]): | ||
result[:, i] = func(values[:, i], start, end, min_periods) | ||
return result | ||
|
||
return column_looper |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,3 @@ | ||
from pandas.core._numba.kernels.mean_ import sliding_mean | ||
|
||
__all__ = ["sliding_mean"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,119 @@ | ||
""" | ||
Numba 1D aggregation kernels that can be shared by | ||
* Dataframe / Series | ||
* groupby | ||
* rolling / expanding | ||
Mirrors pandas/_libs/window/aggregation.pyx | ||
""" | ||
from __future__ import annotations | ||
|
||
import numba | ||
import numpy as np | ||
|
||
|
||
@numba.jit(nopython=True, nogil=True, parallel=False) | ||
def is_monotonic_increasing(bounds: np.ndarray) -> bool: | ||
"""Check if int64 values are monotonically increasing.""" | ||
n = len(bounds) | ||
if n < 2: | ||
return True | ||
prev = bounds[0] | ||
for i in range(1, n): | ||
cur = bounds[i] | ||
if cur < prev: | ||
return False | ||
prev = cur | ||
return True | ||
|
||
|
||
@numba.jit(nopython=True, nogil=True, parallel=False) | ||
def add_mean( | ||
val: float, nobs: int, sum_x: float, neg_ct: int, compensation: float | ||
) -> tuple[int, float, int, float]: | ||
if not np.isnan(val): | ||
nobs += 1 | ||
y = val - compensation | ||
t = sum_x + y | ||
compensation = t - sum_x - y | ||
sum_x = t | ||
if val < 0: | ||
neg_ct += 1 | ||
return nobs, sum_x, neg_ct, compensation | ||
|
||
|
||
@numba.jit(nopython=True, nogil=True, parallel=False) | ||
def remove_mean( | ||
val: float, nobs: int, sum_x: float, neg_ct: int, compensation: float | ||
) -> tuple[int, float, int, float]: | ||
if not np.isnan(val): | ||
nobs -= 1 | ||
y = -val - compensation | ||
t = sum_x + y | ||
compensation = t - sum_x - y | ||
sum_x = t | ||
if val < 0: | ||
neg_ct -= 1 | ||
return nobs, sum_x, neg_ct, compensation | ||
|
||
|
||
@numba.jit(nopython=True, nogil=True, parallel=False) | ||
def sliding_mean( | ||
values: np.ndarray, | ||
start: np.ndarray, | ||
end: np.ndarray, | ||
min_periods: int, | ||
) -> np.ndarray: | ||
N = len(start) | ||
nobs = 0 | ||
sum_x = 0.0 | ||
neg_ct = 0 | ||
compensation_add = 0.0 | ||
compensation_remove = 0.0 | ||
|
||
is_monotonic_increasing_bounds = is_monotonic_increasing( | ||
start | ||
) and is_monotonic_increasing(end) | ||
|
||
output = np.empty(N, dtype=np.float64) | ||
|
||
for i in range(N): | ||
s = start[i] | ||
e = end[i] | ||
if i == 0 or not is_monotonic_increasing_bounds: | ||
for j in range(s, e): | ||
val = values[j] | ||
nobs, sum_x, neg_ct, compensation_add = add_mean( | ||
val, nobs, sum_x, neg_ct, compensation_add | ||
) | ||
else: | ||
for j in range(start[i - 1], s): | ||
val = values[j] | ||
nobs, sum_x, neg_ct, compensation_remove = remove_mean( | ||
val, nobs, sum_x, neg_ct, compensation_remove | ||
) | ||
|
||
for j in range(end[i - 1], e): | ||
val = values[j] | ||
nobs, sum_x, neg_ct, compensation_add = add_mean( | ||
val, nobs, sum_x, neg_ct, compensation_add | ||
) | ||
|
||
if nobs >= min_periods and nobs > 0: | ||
result = sum_x / nobs | ||
if neg_ct == 0 and result < 0: | ||
result = 0 | ||
elif neg_ct == nobs and result > 0: | ||
result = 0 | ||
else: | ||
result = np.nan | ||
|
||
output[i] = result | ||
|
||
if not is_monotonic_increasing_bounds: | ||
nobs = 0 | ||
sum_x = 0.0 | ||
neg_ct = 0 | ||
compensation_remove = 0.0 | ||
|
||
return output |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters