Pyspark：将一列中的值与另一列中同一行中的列表进行匹配

Question

我有一个数据框，其中包含以下 2 列： 1. ID 2. list_IDs

我正在尝试创建第三列，如果 ID 存在于同一行的 list_ID 列中，则返回布尔值 True 或 False

我尝试过使用以下内容：

df = sqlContext.createDataFrame([(1, [1, 2, 3,]), (2, [1, 3, 4])], ("ID", "list_IDs"))

df.withColumn("IDmatch", when(col("ID").isin(F.col("list_IDs")), True).otherwise(False)).show()

那是行不通的。但是，如果我要提供一些静态列表来匹配，它当然可以工作。

df.withColumn("IDmatch", when(col("ID").isin([2, 3]), True).otherwise(False)).show()

我可以使用 udf 返回布尔类型，这也有效：

@udf(returnType=BooleanType())
def isinlist(x, y):
    return x in y

但是，如果可能的话，我试图避免在这种情况下使用 UDF，并且我想知道是否可以使用类似于 .isin() 的本机内容来检查行中的 ID 是否存在于 list_ID 的值列表中同一行的列？

Answer 1

方法一：

如果您使用的是 Spark >= 2.4.0。您可以使用内置arrays_overlap函数。该函数接受 2 个数组并检查它们之间的共同元素。

from pyspark.sql.functions import arrays_overlap, array

df.withColumn("IDmatch", arrays_overlap(df.list_IDs, array(df.ID))).show()

输出：

+---+---------+-------+
| ID| list_IDs|IDmatch|
+---+---------+-------+
|  1|[1, 2, 3]|   true|
|  2|[1, 3, 4]|  false|
+---+---------+-------+

您可以在这里阅读更多相关信息，https://spark.apache.org/docs/2.4.0/api/python/pyspark.sql.html#pyspark.sql.functions.arrays_overlap

方法二：

或者，您也可以使用 anudf来获得相同的输出

from pyspark.sql.functions import udf

element_check = udf(lambda elt_list, elt: elt in elt_list)

df.withColumn("IDmatch", element_check(df.list_IDs, df.ID)).show()