cython 函数在 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

由于亨利已经回答了您问题的熊猫部分，让我谈谈性能方面。 我真的不认为这里需要 Cython。 根据经验，尽可能避免在 np.ndarrays 上循环，即使用矢量化操作/函数而不是循环：

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 开销。 相反，使用.shape属性。
• 使用类型化内存视图快速有效地访问 np.ndarray 的底层内存。
• 通过boundscheck 指令禁用索引检查。
• 同样，您可以通过 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)

亨利提出的解决方案产生：

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)

Answer 2

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