如何在 pyspark 的 pandas_udf 函数中使用正则表达式？

Question

我的代码如下：

@pandas_udf(BooleanType())
def is_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for s in iterator:
        res = re.search("1", s)
        yield res != None

df = spark.createDataFrame(pd.DataFrame(["1", "2", "3"], columns=["v"]))
df.select(is_one(df.v)).show()

我遇到错误：

TypeError: expected string or bytes-like object

看来我的函数没有遍历字符串。这是为什么？我如何在 pandas_udf 中使用正则表达式函数？

我尝试了系列到系列的方法，但得到了同样的错误。

Answer 1

您可以在系列中使用 apply 来应用您的正则表达式搜索：

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import BooleanType
import pandas as pd

@pandas_udf(BooleanType())
def is_one(ser: pd.Series) -> pd.Series:
    return ser.apply(lambda s: re.search("1", s) is not None)
    # a neater way:
    # return ser.str.contains("1")

df = spark.createDataFrame(pd.DataFrame(["1", "2", "3"], columns=["v"]))

df.select(is_one(df.v)).show()
+---------+
|is_one(v)|
+---------+
|     true|
|    false|
|    false|
+---------+