Pandas groupby apply 最后卡住

最近又遇到bug了，排除法排除了半天，终于有定论了。

现象：使用groupby().apply()，最后要出结果的时候程序卡住。

错误代码示例：

import pandas as pd
from tqdm import tqdm


def func(df):
    df['mean_sa'] = df['salary'].shift(1)
    return df

data = [[0, 'C', 43, 35]] * 3000 + \
       [[1, 'C', 18, 30]] * 3000 + \
       [[1, 'A', 20, 22]] * 3000

df_ori = pd.DataFrame(data, columns=['idx', 'company', 'salary', 'age'])
df_ori = df_ori.set_index('idx')
tqdm.pandas(desc='Get returns')
df_ori = df_ori.groupby('company').progress_apply(func)

现象解释：数据结构问题。

当输入数据有很多行，然后index和需要groupby的列有很多重复的数值时，就会出现groupby apply最后卡住的现象。

解决方法很简单：

方法一：对需要分组的那一列先排序！！！本例中是第二列，之后就一切正常了。

# 排序前
data = [[0, 'C', 43, 35]] * 3000 + \
       [[1, 'C', 18, 30]] * 3000 + \
       [[1, 'A', 20, 22]] * 3000

# 排序后
data = [[1, 'A', 43, 35]] * 3000 + \
       [[0, 'C', 18, 30]] * 3000 + \
       [[1, 'C', 18, 30]] * 3000

方法二：使用不同的返回值。修改下面三处代码即可。

import pandas as pd
from tqdm import tqdm


def func(df):
    mean_sa = df['salary'].shift(1)  # 这里不一样
    return mean_sa   # 这里不一样

data = [[0, 'C', 43, 35]] * 3000 + \
       [[1, 'C', 18, 30]] * 3000 + \
       [[1, 'A', 20, 22]] * 3000

df_ori = pd.DataFrame(data, columns=['idx', 'company', 'salary', 'age'])
df_ori = df_ori.set_index('idx')
tqdm.pandas(desc='Get returns')
df_new = df_ori.groupby('company').progress_apply(func)  # 这里不一样

以上就是groupby().apply()卡住的解决方案，希望可以帮到你。欢迎点赞+评论。