cython函数returns应用groupby后单个单元格中的所有值

Question

我一直在寻求加快 groupby 操作的处理速度，虽然现在处理速度更快，但生成的数据帧并不是我想要的。

用一些数据制作多索引数据框：

import pandas as pd
import numpy as np
import cython

data = np.round(np.random.randn(4, 3), 1)

df = pd.DataFrame(data, columns=pd.MultiIndex.from_product([['TOP'], ['A', 'B','C']]))
df

我的 cython 函数如下所示：

%load_ext cython
%%cython
import numpy as np

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return result

当我执行这个函数时：

g = df.groupby(level=0, axis=1).apply(lambda x: func4(x.to_numpy()))

我得到一个单元格中包含所有值的系列：

我想得到的是像这样的索引系列：

如果我在函数末尾将结果转换为数据帧，我可以获得我想要的结果，但是 groupby + apply 再次变慢：

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return pd.DataFrame(result)

编辑：

您可以使用它来生成用于测试的大型数据框：

import pandas as pd
import itertools
import numpy as np
import string
import cython

col = [f'{a}{b}{c}' for a,b,c in list(itertools.product(string.ascii_uppercase[:20], repeat = 3))]
col1 = [str(i) for i in range(1, 10)] # change the 10 to a higher number to produce even more data
col2 = list(i for i in string.ascii_uppercase[:3])

all_keys = ['.'.join(i) for i in itertools.product(col, col1, col2)]
rng = pd.date_range(end=pd.Timestamp.today().date(), periods=6, freq='M')
df = pd.DataFrame(np.random.randint(0, 1000, size=(len(rng), len(all_keys))), columns=all_keys, index=rng)

def top_key(x):
    split = x.split('.', 2)
    return f'{split[0]}.{split[1]}'

df.columns = pd.MultiIndex.from_tuples([(top_key(c), c) for c in df.columns])

print(len(df.columns.get_level_values(0).unique()))
df.head(5)

更新 2，解决方案的性能：

最快，但生成的系列存在上述问题。

使用正确的结果数据帧最快。

Answer 1

groupby.apply 是一个棘手的函数，因为它可以生成聚合值和非聚合值（而且它并不总是选择正确）。

这是一个具体的集合，因此应使用 groupby.aggregate：

def func4(x):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return result


g = df.groupby(level=0, axis=1).aggregate(lambda x: func4(x.to_numpy()))

g:

   TOP
0  1.7
1  2.0
2  0.5
3 -1.1

（可使用 np.random.seed(5) 重现）

import numpy as np
import pandas as pd

np.random.seed(5)
data = np.round(np.random.randn(4, 3), 1)

df = pd.DataFrame(
    data,
    columns=pd.MultiIndex.from_product([['TOP'], ['A', 'B', 'C']])
)

较慢的方法（真的不要这样做）是从 x 传递索引并重建 Series:

def func4(x, idx):
    result = np.zeros([len(x)])
    for i in range(len(x)):
        result[i] = x[i][2] + x[i][1] - x[i][0]
    return pd.Series(result, index=idx)


g = df.groupby(level=0, axis=1).apply(lambda x: func4(x.to_numpy(), x.index))

g:

   TOP
0  1.7
1  2.0
2  0.5
3 -1.1

Answer 2

既然 Henry 已经回答了您问题的 pandas 部分，那么让我谈谈性能方面的问题。我真的不认为这里需要 Cython。根据经验，尽可能避免在 np.ndarrays 上循环，即使用矢量化 operations/functions 而不是循环：

def func4_py_no_loop(x):
    result = np.zeros(x.shape[0])
    result = x[:, 2] + x[:, 1] - x[:, 0]
    return result

但是，如果您真的想使用 Cython，这里有几个关键点可以进一步提高您的初始函数的性能：

键入循环变量 i 以减少 python 开销。
不要多次使用 len。这是一个 python 函数，因此有 python 开销。相反，使用 np.ndarray.

.shape

使用 typed memoryviews 快速高效地访问 np.ndarray 的底层内存。
通过boundscheck directive禁用索引检查。
同样，您可以通过 wraparound 指令禁用对负索引的检查

%%cython

cimport numpy as np
import numpy as np
from cython cimport boundscheck, wraparound

@wraparound(False)
@boundscheck(False)
def func4_cy_loop(double[:, ::1] x):
    cdef int i, N = x.shape[0]
    cdef double[::1] result = np.zeros(N)
    for i in range(N):
        result[i] = x[i, 2] + x[i, 1] - x[i, 0]
    return result

为

计时所有功能

data = np.round(np.random.randn(20000, 3), 1)

Henry 提出的解决方案产生：

In [10]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_cy_your_version(x.to_numpy()))
8.13 ms ± 160 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [11]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_py_no_loop(x.to_numpy()))
1.18 ms ± 24.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [12]: %timeit df.groupby(level=0, axis=1).aggregate(lambda x: func4_cy_loop(x.to_numpy()))
1.3 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

cython函数returns应用groupby后单个单元格中的所有值

cython function returns all values in a single cell after groupby apply

pandas

pandas-groupby

cython