Spark DataFrame: добавьте новые столбцы в соответствии с другими столбцами

Question

Я хочу добавить новый столбец new_col, если значение столбца a находится в yes_list, тогда значение 1 в new_col иначе 0

from pyspark import SparkContext
sc = SparkContext.getOrCreate()
rdd = sc.parallelize([{"a":'y'}, {"a":'y', "b":2}, {"a":'n', "c":3}])
rdd_df = sqlContext.read.json(rdd)

yes_list = ['y']

Примерно так:

rdd_df.withColumn("new_col", [1 if val in yes_list else 0 for val in rdd_df["a"]])

Но приведенное выше неверно и вызывает ошибки.

TypeError: Column is not iterable

Как этого добиться?

Answer 1

Вы можете использовать функции when и isin для API Spark SQL. Это будет go следующим образом:

from pyspark.sql import functions
rdd_df.withColumn("new_col", functions.when(rdd_df['a'].isin(yes_list), 1).otherwise(0)).show()

+---+----+----+-------+                                                         
|  a|   b|   c|new_col|
+---+----+----+-------+
|  y|null|null|      1|
|  y|   2|null|      1|
|  n|null|   3|      0|
+---+----+----+-------+