python之如何根据 Pyspark 中数组列中的值创建新列

我有以下数据框，其中包含代表产品的代码:

testdata = [(0, ['a','b','d']), (1, ['c']), (2, ['d','e'])] 
df = spark.createDataFrame(testdata, ['id', 'codes']) 
df.show() 
+---+---------+ 
| id|    codes| 
+---+---------+ 
|  0|[a, b, d]| 
|  1|      [c]| 
|  2|   [d, e]| 
+---+---------+ 

假设代码 a 和 b 代表 T 恤，代码 c 代表毛衣。

tshirts = ['a','b'] 
sweaters = ['c'] 

如何创建列 label 来检查这些代码是否在数组列中并返回产品名称。像这样:

+---+---------+--------+ 
| id|    codes|   label| 
+---+---------+--------+ 
|  0|[a, b, d]| tshirts| 
|  1|      [c]|sweaters| 
|  2|   [d, e]|    none| 
+---+---------+--------+ 

我已经尝试了很多方法，其中以下是行不通的:

codes = { 
    'tshirts': ['a','b'], 
    'sweaters': ['c'] 
} 
 
def any_isin(ref_values, array_to_search): 
    for key, values in ref_values.items(): 
        if any(item in array_to_search for item in values): 
            return key 
        else: 
            return 'none' 
 
any_isin_udf = lambda ref_values: (F.udf(lambda array_to_search: any_isin_mod(ref_values, array_to_search), StringType())) 
 
df_labeled = df.withColumn('label', any_isin_udf(codes)(F.col('codes'))) 
 
df_labeled.show() 
+---+---------+-------+ 
| id|    codes|  label| 
+---+---------+-------+ 
|  0|[a, b, d]|tshirts| 
|  1|      [c]|   none| 
|  2|   [d, e]|   none| 
+---+---------+-------+ 

请您参考如下方法:

我会用array_contains 来表达。让我们将输入定义为 dict:

from pyspark.sql.functions import expr, lit, when 
from operator import and_ 
from functools import reduce 
 
label_map = {"tshirts": ["a", "b"], "sweaters": ["c"]} 

下一步生成表达式:

expression_map = { 
   label: reduce(and_, [expr("array_contains(codes, '{}')".format(code)) 
   for code in codes]) for label, codes in label_map.items() 
} 

最后用 CASE ... WHEN 减少它:

label = reduce( 
    lambda acc, kv: when(kv[1], lit(kv[0])).otherwise(acc), 
    expression_map.items(),  
    lit(None).cast("string") 
).alias("label") 

结果:

df.withColumn("label", label).show() 
# +---+---------+--------+                                                         
# | id|    codes|   label| 
# +---+---------+--------+ 
# |  0|[a, b, d]| tshirts| 
# |  1|      [c]|sweaters| 
# |  2|   [d, e]|    null| 
# +---+---------+--------+