Как получить pandas groupby-apply для возврата индексов, когда apply возвращает фрейм данных?

Question

При использовании groupby-apply в пандах индексация выходных данных зависит от того, возвращает ли функция apply значение pd.DataFrame или pd.Series.В идеале мне бы хотелось, чтобы моя функция apply возвращала pd.DataFrame, а результат применения должен быть многоиндексным, как это происходит, когда результатом является pd.Series, есть ли способ достичь этого?

Минимальный пример:

import numpy as np
import pandas as pd

data = np.array([np.random.randint(1, 3, 10, dtype=int),
                 np.random.randint(1, 4, 10, dtype=int),
                 np.random.random(10),
                 np.random.random(10)]).T

df = pd.DataFrame(data=data, columns=['A', 'B', 'x', 'y'])
grouped = df.groupby(['A', 'B'])

print(grouped.apply(lambda g: g['x'].diff()))
print(grouped.apply(lambda g: g[['x', 'y']].diff()))

С выводом:

A    B     
1.0  1.0  5         NaN
          8   -0.093603
     2.0  1         NaN
     3.0  2         NaN
          9   -0.399529
2.0  1.0  3         NaN
     2.0  0         NaN
          7   -0.510983
     3.0  4         NaN
          6    0.175588
Name: x, dtype: float64
          x         y
0       NaN       NaN
1       NaN       NaN
2       NaN       NaN
3       NaN       NaN
4       NaN       NaN
5       NaN       NaN
6  0.175588  0.155579
7 -0.510983  0.483819
8 -0.093603  0.465513
9 -0.399529 -0.203326

Я могу придумать несколько обходных путей, но я бы предпочел услышать о каком-то способе сказать, применять для поддержанияиндексирование.

РЕДАКТИРОВАТЬ: Обходной путь, который может быть не так плохо, как я думал, это будет:

pd.concat([grouped.apply(lambda g: g[col].diff()) for col in ['x', 'y']], axis=1)