Почему панды медленнее моей?

Question

У меня есть фрейм данных

            ID  CAT    SCORE
0            0    0  8325804
1            0    1  1484405
...        ...  ...      ...
1999980  99999    0  4614037
1999981  99999    1  1818470

Где я группирую данные по ID и хочу знать 2 категории для каждого идентификатора с наибольшим количеством очков. Я вижу два решения для этого:

df2 = df.groupby('ID').apply(lambda g: g.nlargest(2, columns='SCORE'))

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

Почему ручное управление самое быстрое, чем решение для панд?

MVCE

import numpy as np
import pandas as pd
import time


def create_df(n=10**5, categories=20):
    np.random.seed(0)
    df = pd.DataFrame({'ID': [id_ for id_ in range(n) for c in range(categories)],
                       'CAT': [c for id_ in range(n) for c in range(categories)],
                       'SCORE': np.random.randint(10**7, size=n * categories)})
    return df


def are_dfs_equal(df1, df2):
    columns = sorted(df1.columns)
    if len(df1.columns) != len(df2.columns):
        return False
    elif not all(el1 == el2 for el1, el2 in zip(columns, sorted(df2.columns))):
        return False
    df1_list = [tuple(x) for x in df1[columns].values]
    df1_list = sorted(df1_list, reverse=True)
    df2_list = [tuple(x) for x in df2[columns].values]
    df2_list = sorted(df2_list, reverse=True)
    is_same = df1_list == df2_list
    return is_same


def manual_nlargest(df, n=2):
    df_list = [tuple(x) for x in df[['ID', 'SCORE', 'CAT']].values]
    df_list = sorted(df_list, reverse=True)
    l = []
    current_id = None
    current_id_count = 0
    for el in df_list:
        if el[0] != current_id:
            current_id = el[0]
            current_id_count = 1
        else:
            current_id_count += 1
        if current_id_count <= n:
            l.append(el)
    df = pd.DataFrame(l, columns=['ID', 'SCORE', 'CAT'])
    return df

df = create_df()

t0 = time.time()
df2 = df.groupby('ID').apply(lambda g: g.nlargest(2, columns='SCORE'))
t1 = time.time()
print('nlargest solution: {:0.2f}s'.format(t1 - t0))

t0 = time.time()
df3 = manual_nlargest(df, n=2)
t1 = time.time()
print('manual nlargest solution: {:0.2f}s'.format(t1 - t0))
print('is_same: {}'.format(are_dfs_equal(df2, df3)))

дает

nlargest solution: 97.76s
manual nlargest solution: 4.62s
is_same: True

Answer 1

Я думаю, вы можете использовать это:

df.sort_values(by=['SCORE'],ascending=False).groupby('ID').head(2)

Это то же самое, что и ваше ручное решение с использованием функций сортировки / головы в pandas groupby.

t0 = time.time()
df4 = df.sort_values(by=['SCORE'],ascending=False).groupby('ID').head(2)
t1 = time.time()
df4_list = [tuple(x) for x in df4[['ID', 'SCORE', 'CAT']].values]
df4_list = sorted(df4_list, reverse=True)
is_same = df3_list == df4_list
print('SORT/HEAD solution: {:0.2f}s'.format(t1 - t0))
print(is_same)

дает

SORT/HEAD solution: 0.08s
True

timeit

77.9 ms ± 7.91 ms per loop (mean ± std. dev. of 7 runs, 10 loops each).

Относительно того, почему nlargest медленнее, чем другие решения? Я полагаю, что вызов его для каждой группы создает издержки (%prun показывает 15764409 вызовов функций (15464352 примитивных вызовов) за 30,293 секунды).

Для этого решения (1533 вызова функций (1513 примитивных вызовов) за 0,078 секунды)

Answer 2

Это более быстрое решение, чем ваше ручное решение, если только я не ошибся;) Я думаю, что nlargest () - не самый быстрый способ решения этой проблемы, если вам нужна скорость, но это более читаемое решение.

t0 = time.time()
df4 = df.sort_values(by=['ID', 'SCORE'], ascending=[True, False])
df4['cumcount'] = df4.groupby('ID')['SCORE'].cumcount()
df4 = df4[df4['cumcount'] < 2]
t1 = time.time()
print('cumcount solution: {:0.2f}s'.format(t1 - t0))