Seleccione la fila máxima por grupo – problema de rendimiento pandas

La opción más rápida depende no solo de la longitud del DataFrame (en este caso, alrededor de 13M filas) sino también de la cantidad de grupos. A continuación se presentan los gráficos de puntos que comparan varias formas de encontrar el máximo en cada grupo:

Si solo hay unos pocos grupos (grandes) , using_idxmax puede ser la opción más rápida:

Si hay muchos grupos (pequeños) y el DataFrame no es demasiado grande , using_sort_drop puede ser la opción más rápida:

Sin embargo, using_sort_drop en cuenta que al using_sort_drop , using_sort y using_rank comienzan a verse muy rápido, ya que N = len(df) aumenta, su velocidad relativa a las otras opciones desaparece rápidamente. Para N suficientemente grande, using_idxmax convierte en la opción más rápida , incluso si hay muchos grupos.

using_sort_drop , using_sort y using_rank ordena el DataFrame (o grupos dentro del DataFrame). La clasificación es O(N * log(N)) en promedio, mientras que los otros métodos usan operaciones O(N) . Esta es la razón por la que métodos como using_idxmax beats using_sort_drop para using_sort_drop muy grandes.

Tenga en cuenta que los resultados de los puntos de referencia pueden variar por varios motivos, incluidas las especificaciones de la máquina, el sistema operativo y las versiones de software. Por lo tanto, es importante ejecutar puntos de referencia en su propia máquina y con datos de prueba adaptados a su situación.

Según los perfplots anteriores, using_sort_drop puede ser una opción que vale la pena considerar para su DataFrame de 13M filas, especialmente si tiene muchos grupos (pequeños). De lo contrario, sospecho que using_idxmax es la opción más rápida, pero nuevamente, es importante que compruebe los puntos de referencia en su máquina.

Aquí está la configuración que usé para hacer los perfplots :

 import numpy as np import pandas as pd import perfplot def make_df(N): # lots of small groups df = pd.DataFrame(np.random.randint(N//10+1, size=(N, 2)), columns=['Id','delta']) # few large groups # df = pd.DataFrame(np.random.randint(10, size=(N, 2)), columns=['Id','delta']) return df def using_idxmax(df): return df.loc[df.groupby("Id")['delta'].idxmax()] def max_mask(s): i = np.asarray(s).argmax() result = [False]*len(s) result[i] = True return result def using_custom_mask(df): mask = df.groupby("Id")['delta'].transform(max_mask) return df.loc[mask] def using_isin(df): idx = df.groupby("Id")['delta'].idxmax() mask = df.index.isin(idx) return df.loc[mask] def using_sort(df): df = df.sort_values(by=['delta'], ascending=False, kind='mergesort') return df.groupby('Id', as_index=False).first() def using_rank(df): mask = (df.groupby('Id')['delta'].rank(method='first', ascending=False) == 1) return df.loc[mask] def using_sort_drop(df): # Thanks to jezrael # https://stackoverflow.com/questions/50381064/select-the-max-row-per-group-pandas-performance-issue/50389889?noredirect=1#comment87795818_50389889 return df.sort_values(by=['delta'], ascending=False, kind='mergesort').drop_duplicates('Id') def using_apply(df): selected_idx = df.groupby("Id").apply(lambda df: df.delta.argmax()) return df.loc[selected_idx] def check(df1, df2): df1 = df1.sort_values(by=['Id','delta'], kind='mergesort').reset_index(drop=True) df2 = df2.sort_values(by=['Id','delta'], kind='mergesort').reset_index(drop=True) return df1.equals(df2) perfplot.show( setup=make_df, kernels=[using_idxmax, using_custom_mask, using_isin, using_sort, using_rank, using_apply, using_sort_drop], n_range=[2**k for k in range(2, 20)], logx=True, logy=True, xlabel='len(df)', repeat=75, equality_check=check) 

Otra forma de comparar es usar IPython% timeit :

 In [55]: df = make_df(2**20) In [56]: %timeit using_sort_drop(df) 1 loop, best of 3: 403 ms per loop In [57]: %timeit using_rank(df) 1 loop, best of 3: 1.04 s per loop In [58]: %timeit using_idxmax(df) 1 loop, best of 3: 15.8 s per loop 

Usando el jit de Numba

 from numba import njit import numpy as np @njit def nidxmax(bins, k, weights): out = np.zeros(k, np.int64) trk = np.zeros(k) for i, w in enumerate(weights - (weights.min() - 1)): b = bins[i] if w > trk[b]: trk[b] = w out[b] = i return np.sort(out) def with_numba_idxmax(df): f, u = pd.factorize(df.Id) return df.iloc[nidxmax(f, len(u), df.delta.values)] 

Préstamo de @unutbu

 def make_df(N): # lots of small groups df = pd.DataFrame(np.random.randint(N//10+1, size=(N, 2)), columns=['Id','delta']) # few large groups # df = pd.DataFrame(np.random.randint(10, size=(N, 2)), columns=['Id','delta']) return df 

Prime jit

 with_numba_idxmax(make_df(10)); 

Prueba

 df = make_df(2**20) %timeit with_numba_idxmax(df) %timeit using_sort_drop(df) 47.4 ms ± 99.8 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) 194 ms ± 451 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)