Pyspark SQL：基于另一个DataFrame列中是否存在值，创建一个新列

Question

我尝试按照这个答案的方法操作，但我的问题略有不同。

我有两个pyspark数据框df2和bears2。它们都有一个整数变量，我想创建一个布尔变量，如下伪代码所示：

df3 = df2.withColumn("game", (df2.week_id.isin(bears2.week_if), 1,0))

基本上，如果 df2 的值存在于 bears2 相应的列中，则我想要一个 1，否则是 0。

我尝试了其他问题中的 expr() 方法，但无法使其正常工作。它看起来像这样：

new_column_1 = F.expr(
    """IF(df2.week_id IN(bears2.week_if), 1, 0))"""
    )

- mlewis

---

你忘记了 withColumn(...., when(condition, value).otherwise(another_value))。 - vvg

我根据另一个答案尝试了这个。我需要将bears2.week_if转换为列表/向量才能使其工作吗？因为我得到了一个错误。 - mlewis

@mlewis IN 在数据框中不起作用 - 您必须使用 join。我正在更新答案。 - pault

"isin" 函数需要一个列表作为参数。如果你想要检查另一个数据框，你可以在它上面进行连接操作。 - vvg

Answer 1

目前，您无法像这样在pyspark-sql中使用IN。¹相反，您将不得不连接DataFrames。

尝试类似以下内容：

from pyspark.sql.functions import col, when
df3 = df2.withColumn("id", col("week_id")).alias("df2")\
    .join(bears2.withColumn("id", col("week_if")).alias("bears2"), on="id", how="left")\
    .select("df2.*", when(col("bears2.id").isNotNull(), 1).otherwise(0))

要使连接生效，连接键列必须存在于两个数据框中。出于这个原因，我首先调用withColumn("id", ...)将列重命名为相同的值以进行连接。

接下来我们进行LEFT连接以保留df2中的所有列。最后，我们选择df2中的所有列，并使用pyspark.sql.functions.when()创建布尔列。 when()函数的第一个参数是条件。如果它为True，则返回第二个参数。否则，使用otherwise()中的值。

- pault

---

连接的输出看起来不错，但是when语句返回了“TypeError: condition should be a Column”。我应该尝试之前答案中提到的创建列的方法吗？ - mlewis

@mlewis 请尝试更新- 我已将SQL条件更改为isNotNull()。 - pault