如何检查 Pyspark Dataframe 中是不是存在列表交集

Posted

技术标签:

【中文标题】如何检查 Pyspark Dataframe 中是不是存在列表交集【英文标题】:How to check if there is intersection of lists in Pyspark Dataframe如何检查 Pyspark Dataframe 中是否存在列表交集 【发布时间】:2021-03-29 10:21:50 【问题描述】:

我有一个 pyspark 数据帧如下:

import pyspark.sql.functions as F
import pyspark.sql.types as T
from pyspark.sql.functions import udf

schema = T.StructType([  # schema
    T.StructField("id", T.StringType(), True),
    T.StructField("code", T.ArrayType(T.StringType()), True)])
df = spark.createDataFrame(["id": "1", "code": ["a1", "a2","a3","a4"],
                            "id": "2", "code": ["b1","b2"],
                            "id": "3", "code": ["c1","c2","c3"],
                            "id": "4", "code": ["d1", "b3"]],
                           schema=schema)

给出输出

df.show()


| id|            code|
|---|----------------|
|  1|[a1, a2, a3, a4]|
|  2|        [b1, b2]|
|  3|    [c1, c2, c3]|
|  4|        [d1, b3]|

我希望能够通过向函数提供列和列表来过滤行,如果有任何相交则返回 true(使用与 here 的不相交,因为会有很多非命中)

def lst_intersect(data_lst,query_lst):
    return not set(data_lst).isdisjoint(query_lst) 
lst_intersect_udf = F.udf(lambda x,y: lst_intersect(x,y), T.BooleanType())

当我尝试应用它时

query_lst = ['a1','b3']
df = df.withColumn("code_found", lst_intersect_udf(F.col('code'),F.lit(query_lst)))

得到以下错误

Unsupported literal type class java.util.ArrayList [a1, b3]

我可以通过更改函数等来解决它 - 但想知道我在 F.lit(query_lst) 上做错了什么?

【问题讨论】:

【参考方案1】:

lit 只接受单个值,而不接受 Python 列表。例如,您需要使用列表推导传入包含列表中文字值的数组列。

df2 = df.withColumn(
    "code_found", 
    lst_intersect_udf(
        F.col('code'),
        F.array(*[F.lit(i) for i in query_lst])
    )
)

df2.show()
+---+----------------+----------+
| id|            code|code_found|
+---+----------------+----------+
|  1|[a1, a2, a3, a4]|      true|
|  2|        [b1, b2]|     false|
|  3|    [c1, c2, c3]|     false|
|  4|        [d1, b3]|      true|
+---+----------------+----------+

也就是说,如果您的 Spark >= 2.4,您还可以使用 Spark SQL 函数arrays_overlap 来提供更好的性能:

df2 = df.withColumn(
    "code_found", 
    F.arrays_overlap(
        F.col('code'),
        F.array(*[F.lit(i) for i in query_lst])
    )
)

【讨论】:

以上是关于如何检查 Pyspark Dataframe 中是不是存在列表交集的主要内容,如果未能解决你的问题,请参考以下文章

检查 Pyspark Dataframe 中的重复项

Pyspark - 如何检查两条记录中哪一条具有最新日期及其列值?

如何在 jupyter 中像 pandas Dataframe 一样打印 Pyspark Dataframe

Pyspark - 如何将转换后的列与原始 DataFrame 合并?

PySpark:如何删除 DataFrame 中的非数字列?

PySpark - 如何使用连接更新 Dataframe?