根据另一个数据框列值 p​​yspark 设置列状态

Posted 2023-04-15

【中文标题】根据另一个数据框列值 p​​yspark 设置列状态

【英文标题】：Set column status based on another dataframe column value pyspark

【发布时间】：2021-02-15 11:03:05

【问题描述】：

我有两个 pyspark 数据框，如下所示

main_df：

cat1 cat2 cat3
 a    9     e
 b    3     f
 c    11    g
 d    6     h

support_df：

cat value1 value2
cat1   1      2
cat1   3      4
cat1   5      6
cat1   7      8
cat2   1      2
cat2   3      4
cat2   5      6
cat2   7      8

当cat 列为cat2 时，我需要检查main_df: cat2 中的值是否与support_df: value1 或support_df: value2 中的值匹配。

在上面的例子中，当cat == cat2 时，support_df 中有四行。 main_df 的值是 (9, 3, 11, 6)。其中 value1 和 value2 中只有 3 和 6 个匹配

我想要一个如下的结果数据框

cat1 cat2 cat3 cat2_status
 a    9     e    NotMatched
 b    3     f    Matched
 c    11    g    NotMatched
 d    6     h    Matched

有人可以帮助如何使用 pyspark 实现这一点（我们需要编写用户定义的函数还是可以使用 joins 来完成）？

【参考方案1】：

如果support_df 上有任何匹配项，则使用左连接和when 表达式创建新列cat2_status：

from pyspark.sql import functions as F


result = main_df.alias("main").join(
    support_df.alias("supp"),
    (F.col("supp.cat") == "cat2") &
    ((F.col("main.cat2") == F.col("supp.value1")) |
     (F.col("main.cat2") == F.col("supp.value2"))),
    "left"
).select(
    "main.*",
    F.when(
        F.col("supp.cat").isNotNull(), "Matched"
    ).otherwise("NotMatched").alias("cat2_status")
)

result.show()

#+----+----+----+-----------+
#|cat1|cat2|cat3|cat2_status|
#+----+----+----+-----------+
#|   a|   9|   e| NotMatched|
#|   b|   3|   f|    Matched|
#|   c|  11|   g| NotMatched|
#|   d|   6|   h|    Matched|
#+----+----+----+-----------+