Question

我有一个数据框，其中包含以下两列： 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

您应该定义一个带有2个参数的自定义udf

val isInList = spark.udf
.register("isInList", (id: Int, ids: Seq[Int]) => ids.contains(id))

输出：

df.withColumn("IDmatch", isInList($"ID", $"list_IDs")).show()
+---+---------+-------+
| ID| list_IDs|IDmatch|
+---+---------+-------+
|  1|[1, 2, 3]|   true|
|  2|[1, 3, 4]|  false|
+---+---------+-------+

我正在使用Scala语法，但一般方向应该明确。

Answer 2

方法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

了解更多信息

方法2：

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

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()

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

2 个答案: