python - 使用 numpy 和 numba 提高 groupby-transform 操作的性能

Question

我正在尝试创建一个函数，该函数返回多个浮点数据字段的 z 分数，这些数据由同一二维数组中的分类 str/object 数据分组（最初来自 parquet 文件，通过pandas.read_parquet）。

我正在寻找与使用 pandas 的以下操作等效的功能：

def z_score_series(series: pd.Series) -> np.ndarray:
    values = series.to_numpy()
    return (values - values.mean()) / values.std()

df = pd.read_parquet(<filepath>, columns=<my_columns>, read_dictionary=['x', 'y', 'z'])

df.groupby(['x', 'y', 'z'])['a','b','c'].transform(z_score_series)

此操作使用我的约 270,000 行数据需要 43 秒才能完成，对于我希望用来对数据进行分组的每个附加字段需要更长的时间。我希望简单地提高此操作的速度。

我尝试使用 numpy + numba（最新版本）来提高性能。我知道我分组的字段将始终是分类数据，而我希望计算数字（float32）的 z 分数的字段。

from numba import prange

@numba.njit(fastmath=True)
def z_score(x):
    return (x - x.mean())/x.std()

@numba.njit(parallel=True, fastmath=True)
def transform(unq_idx, values):

    out = np.zeros(values.shape)

    for x in prange(np.max(unq_idx)):
        mask = np.where(unq_idx == x)
        for field in prange(values.shape[0]):
            out[field][mask] = z_score(values[field][mask])

    return out

def get_z_score(df, groupby_cols, calc_cols):

    coded_arr = np.array([df[x].values.codes for x in groupby_cols])
    _, unq_idx = np.unique(coded_arr, axis=1, return_inverse=True)

    calc_arr = np.array([df[x].to_numpy() for x in calc_cols])

    out = transform(unq_idx, calc_arr)
    return out

df = pd.read_parquet(<filepath>, columns=<my_columns>, read_dictionary=['x', 'y', 'z'])
out = get_z_score(df, ['x', 'y', 'z'], ['a', 'b', 'c'])

The variables in the ```get_z_score``` function just for reference.
coded_arr:
array([[  0,   0,   0, ..., 168, 168, 168],
       [  0,   0,   0, ...,   0,   0,   0],
       [  0,   1,   2, ...,   4,   0,   0]], dtype=int16)

unq_idx:
array([    0,     1,     2, ..., 30488, 30484, 30484], dtype=int64)

calc_arr:
array([[ 2.8231514e+07,  6.4926816e+07,  1.3972318e+07, ...,
         1.8851726e+08,  2.0333277e+08,  2.5717517e+08],
       [ 2.9859206e+07,  4.7409168e+07,  1.2298259e+07, ...,
         2.8027318e+08,  1.9047000e+08,  2.3501864e+08],
       [-6.1378721e-02,  7.5533399e-03,  2.3085000e-02, ...,
                   nan,            nan,            nan]], dtype=float32)

使用上面的代码，当增加分组字段的数量时，我可以将等效操作从 43 秒减少到约 3.5 秒，从而显着提高性能。但是，我不相信这种实现是最有效的，并且使用我没有想到的 numpy 函数/技术可以实现更高的效率（我对 numpy 比较陌生）。

分析代码表明这是用于创建输入 z-score 函数的一维向量的索引的计算，它占用了 90% 的代码运行时间，np.where(unq_idx == x)所以这是我认为最大的改进可能是制成。

我尝试用一​​个简单的布尔掩码替换它：mask = unq_idx == x，但这又是一个慢得多的操作（我假设由于必须扫描整个布尔向量来寻找 Trues 而不是直接使用相应的索引np.where）。

所以这个问题可以归结为如何从二维 numpy 数组的一组唯一行中更快地循环、创建和使用索引数组。使用 prange 和 numba 的并行功能会有所帮助，但我想会有一个矢量化或更好的操作可以与使用更多线程相结合。如果您认为这可以通过完全不同的方式（不同的模块？）来实现，或者我的代码可以从原来的位置改进，请告诉我。

谢谢

编辑：这是一些大致相似的随机数据，它们需要大致相似的时间来执行操作。我使用的是 8 个逻辑核心，Intel i7-4770 (3.4GHz)

np.random.seed(0)
df = pd.DataFrame(
    {'x': map(chr, np.random.randint(1,122, 270000)),
     'y': map(chr, np.random.randint(97,97+26, 270000)),
     'z': map(chr, np.random.randint(97,97+10, 270000)),
     'a': np.random.uniform(-1e+8, 1e+8, size=270000),
     'b':np.random.uniform(-1e+8, 1e+8, size=270000),
     'c':np.random.uniform(-1, 1, size=270000)}
)
df[['x','y','z']] = df[['x','y','z']].astype('category')
df[['a','b','c']] = df[['a','b','c']].astype(np.float32)

%timeit get_z_score(df, ['x', 'y', 'z'], ['a', 'b', 'c'])
3.3 s ± 42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df.groupby(['x', 'y', 'z'])[['a', 'b', 'c']].transform(z_score_series)
42.9 s ± 411 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

标签： pythonarrayspandasnumpynumba