Pandas: ¿Cómo hacer que aplicar en el dataframe más rápido?

Considere este ejemplo de pandas en el que estoy calculando la columna C al multiplicar A con B y un float si se cumple una determinada condición utilizando apply con una función lambda :

 import pandas as pd df = pd.DataFrame({'A':[1,2,3,4,5,6,7,8,9],'B':[9,8,7,6,5,4,3,2,1]}) df['C'] = df.apply(lambda x: xA if xB > 5 else 0.1*xA*xB, axis=1) 

El resultado esperado sería:

  ABC 0 1 9 1.0 1 2 8 2.0 2 3 7 3.0 3 4 6 4.0 4 5 5 2.5 5 6 4 2.4 6 7 3 2.1 7 8 2 1.6 8 9 1 0.9 

El problema es que este código es lento y necesito hacer esta operación en un dataframe con alrededor de 56 millones de filas.

El %timeit -result de la operación lambda anterior es:

 1000 loops, best of 3: 1.63 ms per loop 

Al pasar del tiempo de cálculo y también del uso de la memoria al hacer esto en mi gran dataframe, presumo que esta operación utiliza series intermedias al realizar los cálculos.

Traté de formularlo de diferentes maneras, incluyendo el uso de columnas temporales, pero cada solución alternativa que encontré es incluso más lenta.

¿Hay alguna forma de obtener el resultado que necesito de una forma diferente y más rápida, por ejemplo, utilizando numpy ?

Para obtener un mejor rendimiento, es mejor que trabaje con la matriz NumPy y use np.where

 a = df.values # Assuming you have two columns A and B df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1]) 

Prueba de tiempo de ejecución

 def numpy_based(df): a = df.values # Assuming you have two columns A and B df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1]) 

Tiempos –

 In [271]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']]) In [272]: %timeit numpy_based(df) 1000 loops, best of 3: 380 µs per loop In [273]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']]) In [274]: %timeit df['C'] = df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1)) 100 loops, best of 3: 3.39 ms per loop In [275]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']]) In [276]: %timeit df['C'] = np.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B']) 1000 loops, best of 3: 1.12 ms per loop In [277]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']]) In [278]: %timeit df['C'] = np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1)) 1000 loops, best of 3: 1.19 ms per loop 

Mirada más cercana

Echemos un vistazo más de cerca a la capacidad de procesamiento de números de NumPy y comparémosla con pandas en la mezcla –

 # Extract out as array (its a view, so not really expensive # .. as compared to the later computations themselves) In [291]: a = df.values In [296]: %timeit df.values 10000 loops, best of 3: 107 µs per loop 

Caso nº 1: trabaje con la matriz NumPy y use numpy.where:

 In [292]: %timeit np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1]) 10000 loops, best of 3: 86.5 µs per loop 

Nuevamente, asignarlo a una nueva columna: df['C'] tampoco sería muy costoso –

 In [300]: %timeit df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1]) 1000 loops, best of 3: 323 µs per loop 

Caso # 2: trabaje con el dataframe de pandas y use su método .where (no NumPy)

 In [293]: %timeit df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1)) 100 loops, best of 3: 3.4 ms per loop 

Caso 3: trabaje con el dataframe de pandas (sin matriz NumPy), pero use numpy.where

 In [294]: %timeit np.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B']) 1000 loops, best of 3: 764 µs per loop 

Caso n. ° 4: vuelva a trabajar con el dataframe de pandas (sin matriz NumPy), pero use numpy.where

 In [295]: %timeit np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1)) 1000 loops, best of 3: 830 µs per loop 

pandas puros
utilizando pd.Series.where

 df['C'] = df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1)) ABC 0 1 9 1.0 1 2 8 2.0 2 3 7 3.0 3 4 6 4.0 4 5 5 2.5 5 6 4 2.4 6 7 3 2.1 7 8 2 1.6 8 9 1 0.9 

Usando numpy.where :

 df['C'] = numpy.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B']) 

Utilizar:

 df['C'] = np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1)) print (df) ABC 0 1 9 1.0 1 2 8 2.0 2 3 7 3.0 3 4 6 4.0 4 5 5 2.5 5 6 4 2.4 6 7 3 2.1 7 8 2 1.6 8 9 1 0.9