匹配 pyspark dataframe 列以列出并创建新列

Question

我有以下列表。

lst=['name','age','country']

我有以下pyspark dataframe

column_a   column_b
Aaaa       name,age,subject
Bbbb       name,age,country,subject
Cccc       name,subject,percentage

我必须将列表与column_b进行比较，并检查列表中的值是否是列的一部分，然后创建一个新列并使用column_b中可用的列表中的值填充它。

下面是预期的 output。

column_a column_b                 column_c              
Aaaa     name,age,subject         name,age
Bbbb     name,age,country,subject name,age,country
Cccc     name,subject,percentage  name

Answer 1

没有重复

array_intersect允许您想要实现的操作。

array_intersect不允许重复，（即）如果column_b的值为["name", "name"]则column_c将包含一次["name"] 。

from pyspark.sql import functions as F

data = [("Aaaa", ["name", "age", "subject"],),
        ("Bbbb", ["name", "age", "country", "subject"],),
        ("Cccc", ["name", "subject", "percentage"],),
        ("Dddd", ["name", "name"],),]

df = spark.createDataFrame(data, ("column_a", "column_b",))

lst=['name','age','country']
lit_lst = [F.lit(v) for v in lst]

df.withColumn("column_c", F.array_intersect(F.col("column_b"), F.array(lit_lst))).show(truncate=False)

Output

+--------+-----------------------------+--------------------+
|column_a|column_b                     |column_c            |
+--------+-----------------------------+--------------------+
|Aaaa    |[name, age, subject]         |[name, age]         |
|Bbbb    |[name, age, country, subject]|[name, age, country]|
|Cccc    |[name, subject, percentage]  |[name]              |
|Dddd    |[name, name]                 |[name]              |
+--------+-----------------------------+--------------------+

保留重复项

要保留重复项，可以应用filter高阶 Function。

from pyspark.sql import functions as F

data = [("Aaaa", ["name", "age", "subject"],),
        ("Bbbb", ["name", "age", "country", "subject"],),
        ("Cccc", ["name", "subject", "percentage"],),
        ("Dddd", ["name", "name"],),]

df = spark.createDataFrame(data, ("column_a", "column_b",))

df.withColumn("column_c", F.array(lit_lst))\
  .withColumn("column_c", F.expr("filter(column_b, element -> array_contains(column_c, element))"))\
  .show(truncate=False)

Output

+--------+-----------------------------+--------------------+
|column_a|column_b                     |column_c            |
+--------+-----------------------------+--------------------+
|Aaaa    |[name, age, subject]         |[name, age]         |
|Bbbb    |[name, age, country, subject]|[name, age, country]|
|Cccc    |[name, subject, percentage]  |[name]              |
|Dddd    |[name, name]                 |[name, name]        |
+--------+-----------------------------+--------------------+