我能在Pandas身上表演动态累加吗？

Question

如果我有以下数据帧，派生如下：df = pd.DataFrame(np.random.randint(0, 10, size=(10, 1)))

    0
0   0
1   2
2   8
3   1
4   0
5   0
6   7
7   0
8   2
9   2

有没有一种有效的方法cumsum行有一个限制，并且每次达到这个限制时，就开始一个新的cumsum。在达到每个限制（不管行数是多少）之后，将创建一个包含总累计数的行。在

下面我创建了一个这样做的函数的例子，但是它非常慢，特别是当数据帧变得非常大的时候。 我不喜欢我的函数是循环的，我正在寻找一种方法使它更快（我猜是一种没有循环的方法）。在

^{pr2}$

如果你像这样使用我的函数：foo(df, 5) 在上述上下文中，它返回：

   0
2  10
6  8

Answer 1

循环无法避免，但可以使用numba的njit并行化：

from numba import njit, prange

@njit
def dynamic_cumsum(seq, index, max_value):
    cumsum = []
    running = 0
    for i in prange(len(seq)):
        if running > max_value:
            cumsum.append([index[i], running])
            running = 0
        running += seq[i] 
    cumsum.append([index[-1], running])

    return cumsum

这里需要索引，假设您的索引不是数字/单调递增的。在

^{pr2}$

---

如果索引是Int64Index类型，可以将其缩短为：

@njit
def dynamic_cumsum2(seq, max_value):
    cumsum = []
    running = 0
    for i in prange(len(seq)):
        if running > max_value:
            cumsum.append([i, running])
            running = 0
        running += seq[i] 
    cumsum.append([i, running])

    return cumsum

lst = dynamic_cumsum2(df.iloc(axis=1)[0].values, 5)
pd.DataFrame(lst, columns=['A', 'B']).set_index('A')

    B
A    
3  10
7   8
9   4

%timeit foo(df, 5)
1.23 ms ± 30.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit dynamic_cumsum2(df.iloc(axis=1)[0].values, 5)
71.4 µs ± 1.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

---

njit函数性能

perfplot.show(
    setup=lambda n: pd.DataFrame(np.random.randint(0, 10, size=(n, 1))),
    kernels=[
        lambda df: list(cumsum_limit_nb(df.iloc[:, 0].values, 5)),
        lambda df: dynamic_cumsum2(df.iloc[:, 0].values, 5)
    ],
    labels=['cumsum_limit_nb', 'dynamic_cumsum2'],
    n_range=[2**k for k in range(0, 17)],
    xlabel='N',
    logx=True,
    logy=True,
    equality_check=None # TODO - update when @jpp adds in the final `yield`
)

log-log图显示，对于较大的输入，generator函数更快：

一种可能的解释是，随着N的增加，在dynamic_cumsum2中向一个不断增长的列表追加内容的开销变得突出。而cumsum_limit_nb只需要yield。在