python - groupby.shift 时的性能问题

Question

测试代码:

SIZE_MULT = 5
data = np.random.randint(0, 255, size=10**SIZE_MULT, dtype='uint8')
index = pd.MultiIndex.from_product(
            [list(range(10**(SIZE_MULT-1))), list('ABCDEFGHIJ')],
            names = ['d', 'l'])        
test = pd.DataFrame(data, index, columns = ['data'])
test.head()
test['data'].dtype

输出:

        data
d   l   
0   A   137
    B   156
    C   48
    D   186
    E   170

dtype('uint8')

假设我们想要按索引的 0 级分组并移动每个组(例如，移动步长 = 2)。

%%time
shifted = test.groupby(axis=0, level=[0]).shift(2)
print(shifted['data'].dtype)

输出:

float64
CPU times: user 9.43 ms, sys: 56 µs, total: 9.49 ms
Wall time: 8.29 ms

现在问题来了:如果我们想保留我们的 dtype 'uint8'，我们必须摆脱 None，例如，将我们的填充值设置为 0。但是我们现在将获得大量的代码执行时间:

%%time
shifted = test.groupby(axis=0, level=[0]).shift(2, fill_value = 0)
shifted.head()
print(shifted['data'].dtype)

输出:

uint8
CPU times: user 5.9 s, sys: 38.4 ms, total: 5.94 s
Wall time: 5.89 s

所以问题是为什么这么长？如果我们采用没有 fill_value 的第一个移位数据帧，并添加几行代码来获得相同的结果:

%%time
shifted = test.groupby(axis=0, level=[0]).shift(2)
shifted.fillna(0, inplace=True)
shifted = shifted.astype(np.uint8)
print(shifted['data'].dtype)

输出:

uint8
CPU times: user 9.64 ms, sys: 3.68 ms, total: 13.3 ms
Wall time: 11.3 ms

它只会增加几毫秒，而不是 5 秒。

编辑:对应的github issue

Answer 1

来自source code问题是指定的填充值使用缓慢的应用调用。没有填充值，它能够使用更快的 cythonized 结果:

链接中的代码:

def shift(self, periods=1, freq=None, axis=0, fill_value=None):
    #...
    if freq is not None or axis != 0 or not isna(fill_value):
        return self.apply(lambda x: x.shift(periods, freq,
                                            axis, fill_value))

    return self._get_cythonized_result('group_shift_indexer',
                                       self.grouper, cython_dtype=np.int64,
                                       needs_ngroups=True,
                                       result_is_index=True,
                                       periods=periods)

所以在这种情况下，我会在之后使用 .fillna()。