Pyspark; UDF that checks if a column contains one of elements of a list

user_9121931 0
Viewed 6612

I have a dataframe and I want to check if on of its columns contains at least one keywords:

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

When I run this code, I get the following error:

Py4JError: An error occurred while calling z:org.apache.spark.sql.functions.col.
Trace: py4j.Py4JException:
Method col([class java.util.ArrayList]) does not exist

What am I doing wrong?

pyspark user-defined-functions
3 Answers

To make your function work , you should create an array column to compare:

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]     |
+----------------------------------+

So you code would look like below:

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

Outputs:

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

However you could also use isin:

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

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

Another solution that works as well is rlike function. In fact it works much faster than an 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()
user_91219310

The solution that worked is passing the list as a default value, but I still do not understand why the other way did not work:

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