Extracting all matches from different pyspark columns depending on some condition

user_5913102 0
Viewed 1719

I have to extract some codes from columns of a dataframe that looks like the following:

+---------+--------------------------------+--------------------+------+
|first    |second                          |third               |num   |
+---------+--------------------------------+--------------------+------+
|AB12a    |xxxxxx                          |some other data     |100000|
|yyyyyyy  |XYZ02, but possibly also GFH11b |Look at second col* |120000|
+---------+--------------------------------+--------------------+------+

The codes follow the regex "^([A-Z]+[0-9]+[a-z]*)" and are scattered across two columns (first and second) depending on whether the third column contains an asterisk. Since there can be more than one code in each column, I need all the regex matches in an array. In the example above, I need to extract AB12a from first, and [XYZ02, GFH11b] from second.

I found out that multiple matches are not supported by the default pyspark function regexp_extract (https://issues.apache.org/jira/browse/SPARK-24884), so I defined my own regexp_extract_all UDF:

from pyspark.sql.types import *
from pyspark.sql.functions import *
import re

def regexp_extract_all(s, pattern):
    pattern = re.compile(pattern, re.M)
    all_matches = re.findall(pattern, s)
    return all_matches

pattern = "^([A-Z]+[0-9]+[a-z]*)"

udf_regexp_extract_all = udf(regexp_extract_all, ArrayType(StringType()))

I managed to get the UDF working if I apply it on each column separately:

# this extracts AB12a from first
df = df.withColumn("code", udf_regexp_extract_all("first", lit(pattern)))

# this extracts [XYZ02, GFH11b] from second
df = df.withColumn("code", udf_regexp_extract_all("second", lit(pattern)))

But I get a TypeError: expected string or buffer when working in a when clause:

# this gives at runtime TypeError: expected string or buffer
df = df.withColumn("code", when(col("third").like("%*%"), 
                           udf_regexp_extract_all("second", lit(pattern)))
                           .otherwise(udf_regexp_extract_all("first", lit(pattern))))

I think I'm probably getting swamped with types at runtime, because something happens in a when clause that needs my UDF to be defined slightly differently.

Any idea?

pyspark apache-spark user-defined-functions python regex
3 Answers

Your code works perfectly fine. Just change the when condition as below

df.withColumn("code1", when(df.third.rlike("\*") == True, 
                           udf_regexp_extract_all("second", lit("([A-Z]+[0-9]+[a-z]*)")))
                           .otherwise(udf_regexp_extract_all("first", lit("([A-Z]+[0-9]+[a-z]*)")))).show(10, False)

+-------+-------------------------------+-------------------+------+-------+
|first  |second                         |third              |num   |code1  |
+-------+-------------------------------+-------------------+------+-------+
|AB12a  |xxxxxx                         |some other data    |100000|[AB12a]|
|yyyyyyy|XYZ02, but possibly also GFH11b|Look at second col*|120000|[XYZ02, GHF11b]|
+-------+-------------------------------+-------------------+------+-------+
user_100638280

Your code is kind of too verbose and can be simplified by some adjustments for better readability:

Method-1: move logic into the udf

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType
import re

pattern = re.compile(r'\b([A-Z]+[0-9]+[a-z]*)\b')

# use s2 if s3 contains '*', otherwise use s1 (including s3 is NULL)
def regexp_extract_all(s1, s2, s3, pattern):
  try:
    return re.findall(pattern, s2 if '*' in (s3 or '') else s1)
  except:
    return []

udf_regexp_extract_all = udf(lambda x,y,z: regexp_extract_all(x, y, z, pattern), ArrayType(StringType()))

df.withColumn("code", udf_regexp_extract_all('first', 'second', 'third')).show()
#+-------+--------------------+-------------------+------+---------------+
#|  first|              second|              third|   num|           code|
#+-------+--------------------+-------------------+------+---------------+
#|  AB12a|              xxxxxx|    some other data|100000|        [AB12a]|
#|yyyyyyy|XYZ02, but possib...|Look at second col*|120000|[XYZ02, GFH11b]|
#|   null|                 111|                222|  1233|             []|
#+-------+--------------------+-------------------+------+---------------+

Method-2: move when() into udf function argument

from pyspark.sql.functions import udf, when, col

def regexp_extract_all(x, pattern):
    return re.findall(pattern, x or '')

udf_regexp_extract_all = udf(lambda x: regexp_extract_all(x, pattern), ArrayType(StringType()))

df.withColumn('code', udf_regexp_extract_all(when(col('third').like('%*%'), col('second')).otherwise(col('first')))) \
  .show()

Or using Spark SQL expression:

from pyspark.sql.functions import udf, expr

df.withColumn('code', udf_regexp_extract_all(expr("IF(third like '%*%', second, first)"))).show()
user_95107290

In this case, the TypeError happens because there are null values in the columns. I didn't include them in the MWE, but it can (and does!) happen in real data.

An easy way to address this is an additional condition in the regex UDF to make sure re.findall only runs on string objects:

def regexp_extract_all(s, pattern):
    s = "" if s is None else s
    pattern = re.compile(pattern, re.M)
    all_matches = re.findall(pattern, s)
    return all_matches
user_59131020
Related