pandasのgroupbyしてshiftするときの速さ比較

はじめに

pandasでテーブルデータをいじっていると、カテゴリーごとに行方向のシフトを行いたい場合があります（例えば、時系列データで1期ずらしをユーザーごとに行いたいなど）。pandasでグループごとにデータを変換する場合、愚直にはgroupby().transformで実行できます。しかし、グループ化する処理は結構時間がかかります。そこで、いくつかの方法でどれくらい時間がかかるか検証してみます。

環境

• OS 名 Microsoft Windows 10 Home
• プロセッサ Intel(R) Core(TM) i5-7200U CPU @ 2.50GHz、2712 Mhz、2 個のコア、4 個のロジカル プロセッサ
• インストール済みの物理メモリ (RAM) 8.00 GB
• conda 4.8.3
• Python 3.7.7
• numpy==1.18.1
• pandas==1.0.1

手順

データ準備

とりあえず1000万行用意してみました。変数は7つで、5つは適当な数値、2つはカテゴリー変数(X, Y)で、それぞれ10カテゴリー、4カテゴリーとしました。

import numpy as np
import pandas as pd


x = np.arange(10_000_000)
y = np.tile(np.arange(10), int(len(x)/10))
z = np.tile(np.arange(4), int(len(x)/4))

df = pd.DataFrame({"a": x, "b": x, "c": x, "d": x, "e": x, "Y":y, "Z": z})

実験

今回は二つのカテゴリ変数でグループ化してみました。

方法1

愚直な方法で、これが比較の基準になります。

%%timeit -n 1 -r 10

s = df.groupby(["Y", "Z"])["a"].transform(lambda x: x.shift(1))
# 3.25 s ± 107 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

方法2

グループ化する変数をあらかじめ結合しておく方法です。処理は速くなりますが、結合に時間がかかってしまうので、頻繁にグループ化処理をする場合に向いています。

dg = df.copy()
dg["YZ"] = dg["Y"].astype("str") + dg["Z"].astype("str")
# 13.7 s ± 964 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

%%timeit -n 1 -r 10

s = dg.groupby(["YZ"])["a"].transform(lambda x: x.shift(1))
# 2.62 s ± 25.1 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

方法3

pandasのshiftメソッドでなく、numpyのシフト関数を作って実行する方法です。そんなに速さは変わらないようです。

参考：python - Shift elements in a numpy array - Stack Overflow

def shift2(arr, num):
    result = np.empty(arr.shape[0])
    if num > 0:
        result[:num] = np.nan
        result[num:] = arr[:-num]
    elif num < 0:
        result[-num:] = np.nan
        result[:num] = arr[-num:]
    else:
        result = arr
    return result

%%timeit -n 1 -r 10

s = df.groupby(["Y", "Z"])["a"].transform(lambda x: shift2(x, 1))
# 3.2 s ± 15.1 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

方法4

グループ化したあと、イテレーションして各グループを処理する方法です。この方法でも速度がそんなに変わらないのと、より柔軟な処理をかけることができるので、結構好きな方法だったりします。

%%timeit -n 1 -r 10

l = [group["a"].shift(1) for _, group in df.groupby(["Y", "Z"])]    
dh = pd.concat(l, axis=0).sort_index()
# 3.12 s ± 14.4 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

方法5

実はtransform不要です。これが一番早いです。

%%timeit -n 1 -r 10

s = df.groupby(["Y", "Z"])["a"].shift(1)
# 983 ms ± 10.9 ms per loop (mean ± std. dev. of 10 runs, 1 loop each)

結果

方法	説明	time(per loop)
方法１	標準手法	3.25 s ± 0.107 s
方法２	予め結合	2.62 s ± 0.0251 s
方法３	numpy shift	3.2 s ± 0.0151 s
方法４	イテレーション	3.12 s ± 0.0144 s
方法５	transformなし	0.983 s ± 0.0109 s

おわりに

transformを使ってはいけない（戒め）