python - 将 Pandas DataFrame 的 `numpy` 操作转换为 PySpark DataFrame 的等效操作

Question

我正在尝试将 numpy lambda 函数转换为 PySpark 数据框的等效函数。我一直在努力寻找正确的函数/方法。

我对 Pandas 数据框执行以下操作，为 panda_stack 生成一个新列 label:

panda_stack['label'] = panda_stack.groupby('id')['prob1'].transform(lambda x: 0.0 if all(np.isnan(x)) else 1.0)

基本上，我有:

id   prob1
id1  NA
id1  0.12
id2  NA

并且想要:

id  label
id1 1
id2 0

有人可以帮我将上面的代码翻译成对 PySpark 数据框执行相同操作的代码吗？

Answer 1

您可以使用collect_list并应用相同的numpy lambda函数。

import pyspark.sql.functions as F
from pyspark.sql.types import *

check_udf = F.udf(lambda x: 0.0 if all(np.isnan(x)) else 1.0, FloatType())

df = df.groupyby('id').agg(F.collect_list('prob1').alias('label'))\
       .select('id', check_udf(F.col('label').alias('label'))