Могу ли я выполнить динамическое объединение строк в пандах?

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. После того, как достигнут каждый лимит (сколько строк), создается строка с общей суммой.

Ниже я создал пример функции, которая делает это, но она очень медленная, особенно когда кадр данных становится очень большим. Мне не нравится, что моя функция зацикливается, и я ищу способ сделать это быстрее (думаю, без цикла).

def foo(df, max_value):
    last_value = 0
    storage = []
    for index, row in df.iterrows():
        this_value = np.nansum([row[0], last_value])
        if this_value >= max_value:
            storage.append((index, this_value))
            this_value = 0
        last_value = this_value
    return storage

Если вы рассмотрите мою функцию так: foo(df, 5) В приведенном выше контексте он возвращает:

   0
2  10
6  8

Answer 1

Нельзя избежать цикла, но его можно распараллелить, используя numba 's 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

Здесь требуется указатель, при условии, что он не является числовым / монотонно увеличивающимся.

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

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

---

Если индекс имеет тип 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 показывает, что функция генератора быстрее для больших входов:

Возможное объяснение состоит в том, что с увеличением N накладные расходы на добавление в растущий список в dynamic_cumsum2 становятся заметными. В то время как cumsum_limit_nb просто нужно yield.

Answer 2

более простой подход:

def dynamic_cumsum(seq,limit):
    res=[]
    cs=seq.cumsum()
    for i, e in enumerate(cs):
        if cs[i] >limit:
            res.append([i,e])
            cs[i+1:] -= e
    if res[-1][0]==i:
        return res
    res.append([i,e])
    return res

результат:

x=dynamic_cumsum(df[0].values,5)
x
>>[[2, 10], [6, 8], [9, 4]]

Answer 3

Цикл не обязательно плох. Хитрость заключается в том, чтобы убедиться, что он выполняется на объектах низкого уровня. В этом случае вы можете использовать Numba или Cython. Например, используя генератор с numba.njit:

from numba import njit

@njit
def cumsum_limit(A, limit=5):
    count = 0
    for i in range(A.shape[0]):
        count += A[i]
        if count > limit:
            yield i, count
            count = 0

idx, vals = zip(*cumsum_limit(df[0].values))
res = pd.Series(vals, index=idx)

Чтобы продемонстрировать преимущества JIT-компиляции с Numba в производительности:

import pandas as pd, numpy as np
from numba import njit

df = pd.DataFrame({0: [0, 2, 8, 1, 0, 0, 7, 0, 2, 2]})

@njit
def cumsum_limit_nb(A, limit=5):
    count = 0
    for i in range(A.shape[0]):
        count += A[i]
        if count > limit:
            yield i, count
            count = 0

def cumsum_limit(A, limit=5):
    count = 0
    for i in range(A.shape[0]):
        count += A[i]
        if count > limit:
            yield i, count
            count = 0

n = 10**4
df = pd.concat([df]*n, ignore_index=True)

%timeit list(cumsum_limit_nb(df[0].values))  # 4.19 ms ± 90.4 µs per loop
%timeit list(cumsum_limit(df[0].values))     # 58.3 ms ± 194 µs per loop