使用两个字符串参数创建 pandas_udf 函数的最简单方法

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

使用两个字符串参数创建 pandas_udf 函数的最简单方法

The simpliest way to create pandas_udf finction with two string arguments

python

user-defined-functions

pandas

apache-spark

pyspark