基于Pyspark中的列表和列创建列

Question

我有一个pyspark DataFrame，比如说df1，有多列。

我也有一个列表，比如l = ['a','b','c','d']，这些值是DataFrame中某列中存在的值的子集。现在

，我愿做这样的事情：

df2 = df1.withColumn('new_column', expr("case when col_1 in l then 'yes' else 'no' end")) 

但这是抛出以下错误：

failure: "(" expected but identifier l found.

任何想法如何解决此错误或做什么更好的办法它？

Answer 1

你可以做到这一点与Column对象的isin功能：

df1 = sqlContext.createDataFrame([('a', 1), ('b', 2), ('c', 3)], ('col1', 'col2')) 
l = ['a', 'b'] 

from pyspark.sql.functions import * 
df2 = df1.withColumn('new_column', when(col('col1').isin(l), 'yes').otherwise('no')) 

df2.show() 

+----+----+----------+ 
|col1|col2|new_column| 
+----+----+----------+ 
| a| 1|  yes| 
| b| 2|  yes| 
| c| 3|  no| 
+----+----+----------+ 

注：星火< 1.5，使用inSet代替isin。

参考：pyspark.sql.Columndocumentation