python - 根据某些条件从不同的 pyspark 列中提取所有匹配项

Question

我必须从数据框的列中提取一些代码，如下所示:

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

代码遵循正则表达式 "^([A-Z]+[0-9]+[a-z]*)" 并分散在两列中(first 和second) 取决于 third 列是否包含星号。由于每一列中可以有多个代码，因此我需要一个数组中的所有正则表达式匹配项。在上面的示例中，我需要从 first 中提取 AB12a，从 second 中提取 [XYZ02, GFH11b]。

我发现默认的 pyspark 函数 regexp_extract ( https://issues.apache.org/jira/browse/SPARK-24884 ) 不支持多重匹配，所以我定义了自己的 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()))

如果我分别在每一列上应用它，我设法让 UDF 工作:

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

但是在 when 子句中工作时，我得到一个 TypeError: expected string or buffer:

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

我想我可能在运行时被类型淹没了，因为 when 子句中发生了一些事情，需要我的 UDF 的定义略有不同。

有什么想法吗？

Answer 1

您的代码运行良好。只需将 when 条件更改如下

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