Как запустить одну функцию агрегирования для функции pandas.core.groupby.GroupBy, которая создает несколько новых столбцов

Question

Рассмотрим следующий кадр данных:

df = pd.DataFrame({
    'group': [i % 3 for i in range(10)],
    'a': np.random.rand(10),
    'b': np.random.rand(10)
})

def my_agg(x):
    x = x.values.reshape([x.shape[0] // 2,2])
    prod = x[:,0] * x[:,1]
    return [np.sum(prod), np.mean(prod)]

df.set_index('group').stack().groupby('group').apply(my_agg)

Результат как результат

group
0     [0.3625660911145343, 0.09064152277863358]
1       [1.132618561193485, 0.3775395203978283]
2    [0.37300784663400804, 0.12433594887800269]
dtype: object

тогда как я хотел бы иметь отдельный столбец для каждого столбца. Есть ли аккуратный способ сделать это в пандах, учитывая, что:

• несколько созданных функций более сложны, а их совместное вычисление более эффективно;
• количество функций намного больше, чем 2?

Answer 1

Вы можете преобразовать вывод в list с, а затем в DataFrame с помощью конструктора:

def my_agg(x):
    x = x.values.reshape([x.shape[0] // 2,2])
    return [np.sum(x[:,0] * x[:,1]), np.mean(x[:,0] * x[:,1])]

s = df.set_index('group').stack().groupby('group').apply(my_agg)
df1 = pd.DataFrame(s.values.tolist(), index=s.index, columns=['a','b'])
print (df1)
              a         b
group                    
0      2.210601  0.552650
1      0.335913  0.111971
2      1.696796  0.565599

Или вы можете вернуть Series, а затем unstack, но это будет медленнее:

def my_agg(x):
        x = x.values.reshape([x.shape[0] // 2,2])
        return pd.Series([np.sum(x[:,0] * x[:,1]), np.mean(x[:,0] * x[:,1])], index=['a','b'])

df1 = df.set_index('group').stack().groupby('group').apply(my_agg).unstack()
print (df1)
              a         b
group                    
0      0.391921  0.097980
1      0.417366  0.139122
2      0.788845  0.262948