Поведение самого большого (N), когда N больше, чем количество групп?

Question

Я построил DataFrame из следующего списка

df_list_1 = [{"animal": "dog", "color": "red", "age": 4, "n_legs": 4,}, 
             {"animal": "dog", "color": "blue", "age": 4, "n_legs": 3},
             {"animal": "cat", "color": "blue", "age": 4, "n_legs": 4},
             {"animal": "dog", "color": "yellow", "age": 5, "n_legs":2},
             {"animal": "dog", "color": "white", "age": 4, "n_legs": 2},
             {"animal": "dog", "color": "black", "age": 4, "n_legs": 4},
             {"animal": "cat", "color": "brown", "age": 4, "n_legs": 4}]

Теперь я хочу получить новый фрейм данных, который показывает только первые 4 записи (отсортированные по age) на группу с таким же n_legs.

Чтобы сделать это, я попытался

dfg = df_1.set_index(["animal", 'color']).groupby("n_legs")['age'].nlargest(4).reset_index()

Но это дает мне фрейм данных, где столбец n_legs опущен.

    animal  color   age
0   dog     red     4
1   dog     blue    4
2   cat     blue    4
3   dog     yellow  5
4   dog     white   4
5   dog     black   4
6   cat     brown   4

Я думаю, это потому, что 4 равно количеству элементов в самой большой группе.Фактически, если я делаю

dfg = df_1.set_index(["animal", 'color']).groupby("n_legs")['age'].nlargest(3).reset_index()

, я получаю следующее

    n_legs  animal  color   age
0   2       dog     yellow  5
1   2       dog     white   4
2   3       dog     blue    4
3   4       dog     red     4
4   4       cat     blue    4
5   4       dog     black   4

Это предполагаемое поведение?

Есть ли способ всегда отображать столбец даже при использовании nlargest(N) с N большим, чем количество элементов в самой большой группе?

Спасибо!

Answer 1

На мой взгляд, это ошибка 16345 .

Альтернативное решение работает хорошо и, очевидно, быстрее - сначала sort_values, а затем позвоните GroupBy.head:

dfg = (df_1.sort_values(["animal", 'color','age'], ascending=[False, False, True])
          .groupby("n_legs")
          .head(4))