Application de Python function à Pandas grouped DataFrame - quelle est l'approche la plus efficace pour accélérer les calculs?

J'ai affaire à un DataFrame Pandas assez volumineux - mon ensemble de données ressemble à une configuration df suivante:

import pandas as pd
import numpy  as np

#--------------------------------------------- SIZING PARAMETERS :
R1 =                    20        # .repeat( repeats = R1 )
R2 =                    10        # .repeat( repeats = R2 )
R3 =                541680        # .repeat( repeats = [ R3, R4 ] )
R4 =                576720        # .repeat( repeats = [ R3, R4 ] )
T  =                 55920        # .tile( , T)
A1 = np.arange( 0, 2708400, 100 ) # ~ 20x re-used
A2 = np.arange( 0, 2883600, 100 ) # ~ 20x re-used

#--------------------------------------------- DataFrame GENERATION :
df = pd.DataFrame.from_dict(
         { 'measurement_id':        np.repeat( [0, 1], repeats = [ R3, R4 ] ), 
           'time':np.concatenate( [ np.repeat( A1,     repeats = R1 ),
                                    np.repeat( A2,     repeats = R1 ) ] ), 
           'group':        np.tile( np.repeat( [0, 1], repeats = R2 ), T ),
           'object':       np.tile( np.arange( 0, R1 ),                T )
           }
        )

#--------------------------------------------- DataFrame RE-PROCESSING :
df = pd.concat( [ df,
                  df                                                  \
                    .groupby( ['measurement_id', 'time', 'group'] )    \
                    .apply( lambda x: np.random.uniform( 0, 100, 10 ) ) \
                    .explode()                                           \
                    .astype( 'float' )                                    \
                    .to_frame( 'var' )                                     \
                    .reset_index( drop = True )
                  ], axis = 1
                )

Note: Dans le but d'avoir un exemple minimal, il peut être facilement sous-défini (par exemple avec df.loc[df['time'] <= 400, :]), Mais comme je simule les données de toute façon, je pensais que la taille d'origine donnerait un meilleure vue d'ensemble.

Pour chaque groupe défini par ['measurement_id', 'time', 'group'], Je dois appeler la fonction suivante:

from sklearn.cluster import SpectralClustering
from pandarallel     import pandarallel

def cluster( x, index ):
    if len( x ) >= 2:
        data = np.asarray( x )[:, np.newaxis]
        clustering = SpectralClustering( n_clusters   =  5,
                                         random_state = 42
                                         ).fit( data )
        return pd.Series( clustering.labels_ + 1, index = index )
    else:
        return pd.Series( np.nan, index = index )

Pour améliorer les performances, j'ai essayé deux approches:

Paquet Pandarallel

La première approche consistait à paralléliser les calculs en utilisant le package pandarallel:

pandarallel.initialize( progress_bar = True )
df \
  .groupby( ['measurement_id', 'time', 'group'] ) \
  .parallel_apply( lambda x: cluster( x['var'], x['object'] ) )

Cependant, cela semble être sous-optimal car il consomme beaucoup de RAM et tous les cœurs ne sont pas utilisés dans les calculs (même en dépit de la spécification explicite du nombre de cœurs dans la méthode pandarallel.initialize()). Aussi, parfois les calculs se terminent par diverses erreurs, même si je n'ai pas eu la chance de trouver une raison à cela (peut-être un manque de RAM?).

PySpark Pandas UDF

J'ai également essayé un Spark Pandas UDF, bien que je sois totalement nouveau sur Spark. Voici ma tentative:

import findspark;  findspark.init()

from pyspark.sql           import SparkSession
from pyspark.conf          import SparkConf
from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types     import *

spark = SparkSession.builder.master( "local" ).appName( "test" ).config( conf = SparkConf() ).getOrCreate()
df = spark.createDataFrame( df )

@pandas_udf( StructType( [StructField( 'id', IntegerType(), True )] ), functionType = PandasUDFType.GROUPED_MAP )
def cluster( df ):
    if len( df['var'] ) >= 2:
        data = np.asarray( df['var'] )[:, np.newaxis]
        clustering = SpectralClustering( n_clusters   =  5,
                                         random_state = 42
                                         ).fit( data )
        return pd.DataFrame( clustering.labels_ + 1,
                             index = df['object']
                             )
    else:
        return pd.DataFrame( np.nan,
                             index = df['object']
                             )

res = df                                           \
        .groupBy( ['id_half', 'frame', 'team_id'] ) \
        .apply( cluster )                            \
        .toPandas()

Malheureusement, les performances étaient également insatisfaisantes, et d'après ce que j'ai lu sur le sujet, cela peut être juste le fardeau de l'utilisation de la fonction UDF, écrite en Python et le besoin associé de convertir tous les objets Python en Spark objets et inversement.

Voici donc mes questions :

L'une de mes approches pourrait-elle être ajustée pour éliminer les éventuels goulots d'étranglement et améliorer les performances? (par ex. configuration PySpark, ajustement des opérations sous-optimales, etc.)
Y a-t-il de meilleures alternatives? Comment se comparent-ils aux solutions proposées en termes de performances?

pythonpandasapache-sparkparallel-processingdask

24 févr. 2020Kuba_

Je ne suis pas un expert de Dask, mais je fournis le code suivant comme référence:

import dask.dataframe as ddf

df = ddf.from_pandas(df, npartitions=4) # My PC has 4 cores

task = df.groupby(["measurement_id", "time", "group"]).apply(
    lambda x: cluster(x["var"], x["object"]),
    meta=pd.Series(np.nan, index=pd.Series([0, 1, 1, 1])),
)

res = task.compute()

4 mars 2020loopy