Переведите `numpy` для Pandas DataFrame в эквивалентную для PySpark DataFrame

Question

Я пытаюсь перевести лямбда-функцию numpy в ее эквивалент для фрейма данных 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 Data Frame?

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