Как создать новую колонку на pyspark при условии? - PullRequest
Новая
       21

Как создать новую колонку на pyspark при условии?

0 голосов
/ 28 мая 2018

У меня есть следующее data.frame в spark

import findspark
findspark.init()

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
sc = spark.sparkContext
from pyspark.sql import functions as sf
from pyspark.sql.functions import col, when, lit

ddf = spark.createDataFrame([[None, 'Michael',2],
                             [30, 'Andy',3],
                             [19, 'Justin',4],
                             [30, 'James Dr No From Russia with Love Bond',6]],
                            schema=['age', 'name','weights'])
ddf.show()

В этом тривиальном примере я хотел бы создать два столбца: один с weighted.mean из age, если age>29(с именем weighted_age) и другим age^2, если age<=29 (с именем age_squared)

1 Ответ

0 голосов
/ 28 мая 2018

Сначала вы должны найти weighted.mean из всего набора данных с помощью age > 29, а затем заполнить с помощью withColumn.Это потому, что weighted.mean зависит от всего набора данных.

age_squared можно сделать построчно как 

from pyspark.sql import functions as f
weightedMean = ddf.filter(f.col('age')>29).select(f.sum(f.col('age')*f.col('weights'))/f.sum(f.col('weights'))).first()[0]

ddf.withColumn('weighted_age', f.when(f.col('age') > 29, weightedMean))\
    .withColumn('age_squared', f.when(f.col('age') <= 29, f.col('age')*f.col('age')))\
    .show(truncate=False)

, что должно дать вам 

+----+--------------------------------------+-------+------------+-----------+
|age |name                                  |weights|weighted_age|age_squared|
+----+--------------------------------------+-------+------------+-----------+
|null|Michael                               |2      |null        |null       |
|30  |Andy                                  |3      |30.0        |null       |
|19  |Justin                                |4      |null        |361        |
|30  |James Dr No From Russia with Love Bond|6      |30.0        |null       |
+----+--------------------------------------+-------+------------+-----------+

Youможно заполнить другое значение, используя .otherwise с функцией when вместо заполнения по умолчанию null

...