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

【问题标题】：Extracting all matches from different pyspark columns depending on some condition根据某些条件从不同的 pyspark 列中提取所有匹配项
【发布时间】：2020-02-24 13:58:28
【问题描述】：

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

+---------+--------------------------------+--------------------+------+
|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_allUDF：

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 进行稍微不同的定义。

有什么想法吗？

【问题讨论】：

我认为你应该使用 like("%*%") 而不是使用 like("*") Like THIS df = df.withColumn("code", when(col("third") .like("%*%"), udf_regexp_extract_all("second", lit(pattern))) .otherwise(udf_regexp_extract_all("first", lit(pattern))))
你说得对，我打错了。我在问题中对其进行了编辑，但这不是问题的根本原因。
你为我使用 bcoz 的 spark 版本吗 spark 2.4+ 在相同的代码上没有给出任何错误。并尝试从模式中删除 ^ 以将数据作为数组检索。
我正在运行 2.2.0，但问题是其中一列中可能存在空值。
因此，您能否在数据框中也使用 null 值来处理该错误

标签： python regex apache-spark pyspark user-defined-functions

【解决方案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]|
+-------+-------------------------------+-------------------+------+-------+

【讨论】：

非常感谢您的仔细检查！事实上，它确实可以在上面的 MWE 上运行，但不能在我的真实数据上运行。它让我进一步调查，我意识到问题是第一列中可能存在导致 TypeError 的空值。我会将此作为答案发布以提高知名度。

【解决方案2】：

您的代码有点过于冗长，可以通过一些调整来简化以提高可读性：

方法一：将逻辑移入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|             []|
#+-------+--------------------+-------------------+------+---------------+

方法二：将 when() 移入 udf 函数参数

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

或者使用 Spark SQL 表达式：

from pyspark.sql.functions import udf, expr

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

【讨论】：

我一直在寻求代码样式方面的建议——感谢您的意见。

【解决方案3】：

在这种情况下，TypeError 发生是因为列中有空值。我没有将它们包含在 MWE 中，但它可以（并且确实！）发生在真实数据中。

解决此问题的一种简单方法是在正则表达式 UDF 中添加一个附加条件，以确保 re.findall 仅在字符串对象上运行：

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

【讨论】：