使用PySpark DataFrame的groupBy计算AUC(使用pandas_udf定义聚合函数)


介绍

当使用

PySpark计算AUC时,可以使用BinaryClassificationEvaluator类轻松地计算它。
但是,为了理解模型之间的差异,不可能满足为每个段而不是整个测试数据计算AUC的需要。

作为一种解决方法,我定义了一个聚合函数,该函数使用pandas_udf计算AUC,并使用agg方法进行计算。

实施实例

示例如下。

预先计算正确答案标签(true)和预测分数(pred),并通过引用它们来定义计算AUC的聚合函数。

请注意,在

agg方法的pandas_udf中定义的聚合函数不能与spark准备的聚合函数一起使用。
(如果尝试一起使用,则会出现错误Cannot use a mixture of aggregate function and group aggregate pandas UDF)

UDF定义

1
2
3
4
5
6
7
8
from sklearn.metrics import roc_auc_score
from pyspark.sql.types import DoubleType
from pyspark.sql.functions import pandas_udf, PandasUDFType


@pandas_udf(DoubleType(), functionType=PandasUDFType.GROUPED_AGG)
def auc_udf(true, pred):
    return roc_auc_score(true, pred)

计算方法

1
data.groupBy(key).agg(auc_udf('true', 'pred').alias('auc'))

参考

  • 结合使用Spark和Pandas UDF
  • apache spark-在PySpark中的GroupedData上应用UDF(带有可运行的python示例)