【Python】pandasでLAG / LEAD OVER PARTITION BY したいとき【SQL】

SQL書きすぎてるとpandasの書き方わからなくなる。

pandasでSQLの

LAG(value) OVER(PARTITION BY class ORDER BY value)

のような処理をどうやって書くのか？
groupbyとshiftを使えば一発です。LEAD()の場合はshift(-1)でOK。

import pandas as pd

df = pd.DataFrame({
    'class': ['A', 'A' ,'A' ,'B' ,'B' ,'B'],
    'value': [1, 2, 3, 4, 5, 6]}
)

df_lag = pd.DataFrame(df.groupby(['class']).shift(1))
df_lag = df_lag.rename(columns={'value': 'lag_value'})
pd.concat([df, df_lag], axis=1)

> output
class value    lag_value
0  A   1  NaN
1  A   2  1.0
2  A   3  2.0
3  B   4  NaN
4  B   5  4.0
5  B   6  5.0