创建带有两个字符串参数的pandas_udf fintion的最简单方法

Question

我需要使用pandas_udf对基于两列的字符串进行某些计算。

我的代码如下：

@pandas_udf(BooleanType())
def is_one_pd(ser: pd.Series, ser2: pd.Series) -> pd.Series:
def is_one(s):
return (re.search("1", s.title) is not None) & (re.search("1", s.comp_short) is not None)
df = pd.DataFrame({'title':ser, 'comp_short':ser2})
return df.apply(is_one, 1)
df = spark.createDataFrame(pd.DataFrame([("1", "1"), ("2", "2"), ("3", "3")], columns=["title", "comp_short"]))
df.select(is_one_pd(df.title, df.comp_short)).show()

它看起来太复杂了。还有更好的选择吗？我需要使用pandas_udf，因为实际的示例更复杂。

Answer 1

您可以直接在系列上应用正则表达式搜索，而无需构建数据帧：

import re 
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import BooleanType
import pandas as pd
@pandas_udf(BooleanType())
def is_one_pd(ser: pd.Series, ser2: pd.Series) -> pd.Series:
return (ser.apply(lambda s: re.search("1", s) is not None) & 
ser2.apply(lambda s: re.search("1", s) is not None))
df = spark.createDataFrame(pd.DataFrame([("1", "1"), ("2", "2"), ("3", "3")], columns=["title", "comp_short"]))
df.select(is_one_pd(df.title, df.comp_short)).show()