zmc
2023-08-08 e792e9a60d958b93aef96050644f369feb25d61b
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
"""
datetimelke_accumulations.py is for accumulations of datetimelike extension arrays
"""
 
from __future__ import annotations
 
from typing import Callable
 
import numpy as np
 
from pandas._libs import iNaT
 
from pandas.core.dtypes.missing import isna
 
 
def _cum_func(
    func: Callable,
    values: np.ndarray,
    *,
    skipna: bool = True,
):
    """
    Accumulations for 1D datetimelike arrays.
 
    Parameters
    ----------
    func : np.cumsum, np.maximum.accumulate, np.minimum.accumulate
    values : np.ndarray
        Numpy array with the values (can be of any dtype that support the
        operation). Values is changed is modified inplace.
    skipna : bool, default True
        Whether to skip NA.
    """
    try:
        fill_value = {
            np.maximum.accumulate: np.iinfo(np.int64).min,
            np.cumsum: 0,
            np.minimum.accumulate: np.iinfo(np.int64).max,
        }[func]
    except KeyError:
        raise ValueError(f"No accumulation for {func} implemented on BaseMaskedArray")
 
    mask = isna(values)
    y = values.view("i8")
    y[mask] = fill_value
 
    if not skipna:
        mask = np.maximum.accumulate(mask)
 
    result = func(y)
    result[mask] = iNaT
 
    if values.dtype.kind in ["m", "M"]:
        return result.view(values.dtype.base)
    return result
 
 
def cumsum(values: np.ndarray, *, skipna: bool = True) -> np.ndarray:
    return _cum_func(np.cumsum, values, skipna=skipna)
 
 
def cummin(values: np.ndarray, *, skipna: bool = True):
    return _cum_func(np.minimum.accumulate, values, skipna=skipna)
 
 
def cummax(values: np.ndarray, *, skipna: bool = True):
    return _cum_func(np.maximum.accumulate, values, skipna=skipna)