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

Estoy seleccionando una fila máxima por grupo y estoy usando groupby / agg para devolver los valores de índice y seleccionar las filas usando loc .

Por ejemplo, para agrupar por "Id" y luego seleccionar la fila con el valor más alto de "delta" :

 selected_idx = df.groupby("Id").apply(lambda df: df.delta.argmax()) selected_rows = df.loc[selected_idx, :] 

Sin embargo, es tan lento de esta manera. En realidad, mi laptop i7 / 16G RAM se cuelga cuando estoy usando esta consulta en 13 millones de filas.

Tengo dos preguntas para expertos:

  1. ¿Cómo puedo hacer que esta consulta se ejecute rápidamente en pandas? ¿Qué estoy haciendo mal?
  2. ¿Por qué esta operación es tan cara?

[Actualización] ¡Muchas gracias por el análisis de @unutbu! sort_drop lo es! En mi máquina i7 / 32GRAM, groupby + idxmax se cuelga durante casi 14 horas (nunca devuelve nada) sin embargo, sort_drop manejó ¡MENOS DE UN MINUTO!

¡Todavía tengo que ver cómo los pandas implementan cada método, pero los problemas están resueltos por ahora! Me encanta StackOverflow.

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: introduzca la descripción de la imagen aquí

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: introduzca la descripción de la imagen aquí

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)