派斯帕克；用于检查列是不是包含列表元素之一的 UDF

Posted 2023-04-13

【中文标题】派斯帕克；用于检查列是不是包含列表元素之一的 UDF

【英文标题】：Pyspark; UDF that checks if a column contains one of elements of a list派斯帕克；用于检查列是否包含列表元素之一的 UDF

【发布时间】：2020-05-06 13:18:59

【问题描述】：

我有一个数据框，我想检查它的列中是否包含至少一个关键字：

from pyspark.sql import types as T
import pyspark.sql.functions as fn
key_labels = ["COMMISSION", "COM", "PRET", "LOAN"]

def containsAny(string, array):
    if len(string) == 0:
        return False
    else:
        return (any(word in string for word in array))

contains_udf = fn.udf(containsAny, T.BooleanType())

df = spark.createDataFrame([("COMMISSION", "1"), ("CAMMISSION", "2")], ("original", "id"))
df.withColumn("keyword_match", contains_udf(fn.col("original"),key_labels)).show()

当我运行此代码时，我收到以下错误：

Py4JError：调用 z:org.apache.spark.sql.functions.col 时出错。 跟踪：py4j.Py4JException: 方法 col([class java.util.ArrayList]) 不存在

我做错了什么？

【问题讨论】：

这能回答你的问题吗？ Filter pyspark dataframe if contains a list of strings 或者可能是这个***.com/questions/48869922/…

可以，但是我想知道为什么在这种情况下将列表传递给 udf 不起作用？

在您的情况下，udf 期望 key_labels 是一列或一列的名称，而不是。这几乎回答了您的其他问题：***.com/a/47912902/7306659

【参考方案1】：

为了让你的函数工作，你应该创建一个数组列来比较：

df.select(fn.array([fn.lit(i) for i in key_labels])).show(truncate=False)

---

+----------------------------------+
|array(COMMISSION, COM, PRET, LOAN)|
+----------------------------------+
|[COMMISSION, COM, PRET, LOAN]     |
|[COMMISSION, COM, PRET, LOAN]     |
+----------------------------------+

所以你的代码如下所示：

def containsAny(string, array):
    if len(string) == 0:
        return False
    else:
        return (any(word in string for word in array))

contains_udf = fn.udf(containsAny, T.BooleanType())
(df.withColumn("keyword_match", contains_udf(fn.col("original"),
 fn.array([fn.lit(i) for i in key_labels])))).show()

---

输出：

+----------+---+-------------+
|  original| id|keyword_match|
+----------+---+-------------+
|COMMISSION|  1|         true|
|CAMMISSION|  2|        false|
+----------+---+-------------+

不过你也可以使用isin:

df.withColumn('keyword_match',df['original'].isin(key_labels)).show()

+----------+---+-------------+
|  original| id|keyword_match|
+----------+---+-------------+
|COMMISSION|  1|         true|
|CAMMISSION|  2|        false|
+----------+---+-------------+

【讨论】：

没错，不需要 udf，isin 是要走的路

我不使用 isin 的原因是因为 original 包含其他符号。因此，要应用此解决方案，我需要先将字符串拆分为单词，然后循环遍历数组，但有时我要搜索的字符串会同时包含多个单词。所以函数 contains 是正确的选择...

【参考方案2】：

另一个同样有效的解决方案是rlike 函数。事实上，它的运行速度比udf 快得多。

regex = "|".join(r"(" + x + r")" for x in key_labels)
df = spark.createDataFrame([("COMMISSION", "1"), ("CAMMISSION", "2")], ("original", "id"))
df.select("original","id",fn.col("original").rlike(regex).alias("keyword_match")).show()

【参考方案3】：

有效的解决方案是将列表作为默认值传递，但我仍然不明白为什么其他方式不起作用：

def containsAny(string, array=key_labels):
    if len(string) == 0:
        return False
    else:
        return (any(word in string for word in array))

contains_udf = fn.udf(containsAny, T.BooleanType())

df = spark.createDataFrame([("COMMISSION", "1"), ("CAMMISSION", "2")], ("original", "id"))
df.withColumn("keyword_match", contains_udf(fn.col("original"))).show()