将 Pandas DataFrame 的 'numpy' 操作转换为 PySpark DataFrame 的等效操作

Question

我正在尝试将numpylambda函数转换为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'))

Answer 2

我无法给出一个直接的答案。但是在其他人回答之前，您可以使用以下代码

import numpy as np
import pyspark.sql.functions as F
df = sqlContext.createDataFrame([(1, np.nan), (1, 0.12), (2, np.nan)], ('id', 'prob1'))
df = df.withColumn(
'prob1',
F.when(
F.col('prob1') == 0,
F.lit(0.01)
).otherwise(
F.col('prob1')
)
)
df = df.fillna(0)
df = df.groupBy('id').agg(
F.sum(F.col('prob1')).alias('label')
)
df = df.withColumn(
'label',
F.when(
F.col('label') != 0,
F.lit(1)
).otherwise(
F.col('label')
)
)
df.show()