Pyspark; UDF, который проверяет, содержит ли столбец один из элементов списка - PullRequest
Новая
       91

Pyspark; UDF, который проверяет, содержит ли столбец один из элементов списка

1 голос
/ 06 мая 2020

У меня есть фрейм данных, и я хочу проверить, содержит ли один из его столбцов хотя бы одно ключевое слово: 

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.
Trace: py4j.Py4JException:
Method col ([class java .util .ArrayList]) не существует

Что я делаю не так?

Ответы [ 2 ]

2 голосов
/ 06 мая 2020

Чтобы ваша функция работала, вы должны создать столбец массива для сравнения: 

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|
+----------+---+-------------+
0 голосов
/ 06 мая 2020

Сработавшее решение - передача списка в качестве значения по умолчанию, но я все еще не понимаю, почему другой способ не сработал: 

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