python - Pyspark 基于具有列表或集合的多个条件的其他列创建新列

Question

我正在尝试在 pyspark 数据框中创建一个新列。我有如下数据

+------+
|letter|
+------+
|     A|
|     C|
|     A|
|     Z|
|     E|
+------+

我想根据给定列添加新列

+------+-----+
|letter|group|
+------+-----+
|     A|   c1|
|     B|   c1|
|     F|   c2|
|     G|   c2|
|     I|   c3|
+------+-----+

可以有多个类别，有许多单独的字母值(大约100个，也包含多个字母)

我已经用 udf 完成了这个，并且运行良好

from pyspark.sql.functions import udf
from pyspark.sql.types import *

c1 = ['A','B','C','D']
c2 = ['E','F','G','H']
c3 = ['I','J','K','L']
...

def l2c(value):
    if value in c1: return 'c1'
    elif value in c2: return 'c2'
    elif value in c3: return 'c3'
    else: return "na"

udf_l2c = udf(l2c, StringType())
data_with_category = data.withColumn("group", udf_l2c("letter"))

现在我尝试在没有udf的情况下做到这一点。也许使用 when 和 col。我尝试过的如下。它可以工作，但是代码很长。

data_with_category = data.withColumn('group', when(col('letter') == 'A' ,'c1')
    .when(col('letter') == 'B', 'c1')
    .when(col('letter') == 'F', 'c2')
    ... 

为字母的所有可能值编写新的when条件很长而且不太好。在我的例子中，字母的数量可能非常大(大约 100 个)。所以我尝试了

data_with_category = data.withColumn('group', when(col('letter') in ['A','B','C','D'] ,'c1')
    .when(col('letter') in ['E','F','G','H'], 'c2')
    .when(col('letter') in ['I','J','K','L'], 'c3')

但它返回错误。我该如何解决这个问题？

Answer 1

使用isin .

c1 = ['A','B','C','D']
c2 =['E','F','G','H']
c3 =['I','J','K','L']

df.withColumn("group", F.when(F.col("letter").isin(c1),F.lit('c1'))\
                        .when(F.col("letter").isin(c2),F.lit('c2'))\
                        .when(F.col("letter").isin(c3),F.lit('c3'))).show()

#+------+-----+
#|letter|group|
#+------+-----+
#|     A|   c1|
#|     B|   c1|
#|     F|   c2|
#|     G|   c2|
#|     I|   c3|
#+------+-----+