将 Python 函数应用于 Pandas 分组数据帧 - 加速计算的最有效方法是什么？

Question

我正在处理相当大的 Pandas DataFrame - 我的数据集类似于以下df设置：

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
                )

注意：为了有一个最小的例子，它可以很容易地被子集化（例如df.loc[df['time'] <= 400, :]），但由于我无论如何模拟数据，我认为原始大小会提供更好的概览。

对于由['measurement_id', 'time', 'group']我定义的每个组，我需要调用以下函数：

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 )

为了提高性能，我尝试了两种方法：

潘达列包

第一种方法是使用pandarallel包并行计算：

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

然而，这似乎是次优的，因为它消耗大量 RAM 并且并非所有内核都用于计算（即使在 pandarallel.initialize()方法中）。此外，有时计算会因各种错误而终止，尽管我还没有机会找到原因（可能是 RAM 不足？）。

PySpark Pandas UDF

我还尝试了 Spark Pandas UDF，尽管我对 Spark 完全陌生。这是我的尝试：

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()

不幸的是，性能也不令人满意，从我读到的主题来看，这可能只是使用 Python 编写的 UDF 函数的负担以及将所有 Python 对象转换为 Spark 对象并返回的相关需求。

所以这里是我的问题：

1. 是否可以调整我的任何一种方法以消除可能的瓶颈并提高性能？（例如 PySpark 设置、调整次优操作等）
2. 他们有更好的选择吗？在性能方面，它们与提供的解决方案相比如何？

Answer 1

这不是一个答案，但是......

如果你跑

df.groupby(['measurement_id', 'time', 'group']).apply(
    lambda x: cluster(x['var'], x['object']))

（即，仅使用 Pandas），您会注意到您已经在使用多个内核。这是因为默认情况下sklearn使用joblib并行工作。您可以将调度程序换成 Dask，也许会比在线程之间共享数据获得更高的效率，但只要您所做的工作像这样受 CPU 限制，您就无法采取任何措施来加快速度。

简而言之，这是一个算法问题：在尝试考虑不同的计算框架之前，先弄清楚您真正需要计算什么。