根据Pyspark中的列表和列创建列

时间:2016-05-16 14:17:38

标签: python-3.x apache-spark pyspark

我有一个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"))

但这会引发以下错误:

  

失败:“(”预期但是找到了标识符。

知道如何解决此错误或更好的方法吗?

1 个答案:

答案 0 :(得分:3)

您可以使用isin对象的Column函数实现此目的:

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|
+----+----+----------+

注意:对于Spark< 1.5,使用inSet代替isin

参考:pyspark.sql.Column documentation