pandas groupby 中的 apply 函数可以返回多个数据帧吗？

Question

我想在 apply 函数中为 groupby 执行两个略有不同的操作，如下代码所示：

import pandas as pd
df = pd.DataFrame({'x' : [1,2,2], 'y':['aaa','bbb','ccc']})
print(df)

def func1(x):
    stackrow = pd.concat([x, x], axis=0, keys=['a', 'b'], names=['w', 'z'])
    return(stackrow)

df2 = df.groupby('x').apply(func1)
print("***********************\ndf2:")
print(df2)

def func2(x):
    stackcol = pd.concat([x, x], axis=1, keys=['c', 'd'], names=['w', 'z'])
    return(stackcol)

df3 = df.groupby('x').apply(func2)
print("***********************\ndf3:")
print(df3)

但是，正如我在实际示例中提到的 func1 和 func2 具有显着重叠的功能，我不想重新执行该通用功能的代码。我想知道下面的代码是否可以工作。

def func(x):
    stackrow = pd.concat([x, x], axis=0, keys=['a', 'b'], names=['z'])
    stackcol = pd.concat([x, x], axis=1, keys=['c', 'b'], names=['w'])
    return(stackrow, stackcol)

df4, df5 = df.groupby('x').apply(func)

print("***********************\ndf4:")
print(df4)
print("***********************\ndf5:")
print(df5)

但它似乎不适用于df2!=df4 和df3!=df5。有什么想法吗？

Answer 1

使用嵌套的 DataFrame 似乎可行。

def func(df):
    # reseting index to avoid some error in pd.concat later not sure why
    x = df.reset_index(drop=True)
    stackrow = pd.concat([x, x], axis=0, keys=['a', 'b'], names=['w', 'z'])
    stackcol = pd.concat([x, x], axis=1, keys=['c', 'd'], names=['p', 'q'])
    newdf = pd.DataFrame([[stackrow, stackcol]])
    return(newdf)

tmp = df.groupby('x').apply(func)

df4 = pd.concat(tmp[0].tolist())
df5 = pd.concat(tmp[1].tolist())
print("***********************\ndf4:\n", df4)
print("***********************\ndf5:\n", df5)

或者，最好将分组变量保留在索引中：

df44 = tmp[[0]].groupby('x').apply(lambda x: x.iloc[0,0])
df55 = tmp[[1]].groupby('x').apply(lambda x: x.iloc[0,0])
print("***********************\ndf44:\n", df44)
print("***********************\ndf55:\n", df55)

Answer 2

您可以遍历组并直接调用函数

gb = df.groupby('x') 
for group_name, group_lables in gb.groups.items(): 
    df4, df5 = func(gb.get_group(group_name)) 
    # do something more with the dfs?