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

时间:2018-06-28 03:29:40

标签: python pandas dataframe pyspark

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

我对Pandas数据框执行以下操作,以为label生成新列panda_stack

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数据框相同的内容吗?

2 个答案:

答案 0 :(得分: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'))

答案 1 :(得分:0)

我不能直接给出答案。但是直到别人回答您​​可以使用下面的代码

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