如何在PySpark中正确使用Pandas UDF实现波斯日期转公历日期

问题背景 #

有一个百万级记录的PySpark DataFrame，其中persian_date列是8位波斯日期字符串（如14010901），需要转换为公历日期（如2022-11-22）。先后尝试Python UDF（性能差）、Scala UDF（提升不明显），转而尝试Pandas UDF但多次报错，以下是尝试过程及错误：

尝试1：Scalar型UDF报错 #

import pandas as pd

@pandas_udf('long', PandasUDFType.SCALAR)
def f1(v: pd.Series) -> pd.Series:
  return v.map(lambda x: JalaliDate(int(str(x[1])[0:4]), int(str(x[1])[4:6]), int(str(x[1])[6:8])).to_gregorian())

df.withColumn('date_miladi', f1(df.trx_date)).show()

错误：TypeError: 'decimal.Decimal' object is not subscriptable

尝试2：迭代器型UDF报错 #

import pandas as pd
from typing import Iterator
@pandas_udf(DateType())
def f1(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
 for date in iterator:
     return pd.Series(JalaliDate(int(str(date[1])[0:4]), int(str(date[1])[4:6]), int(str(date[1])[6:8])).to_gregorian())

df.withColumn('date_miladi', f1(df.trx_date)).show()

错误：TypeError: Return type of the user-defined function should be Pandas.Series, but is <class 'datetime.date'>

尝试3：迭代器型UDF再次报错 #

@pandas_udf("string",PandasUDFType.SCALAR_ITER)
def f1(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for date in iterator:
        print('type date:', type(date[1]))
        yield str(JalaliDate(int(str(date[1])[0:4]), int(str(date[1])[4:6]), int(str(date[1])[6:8])).to_gregorian())

错误：AttributeError: 'str' object has no attribute 'isnull'

DataFrame样例：

+-----------+-------------+
 |id         | persian_date|
 +-----------+-------------+
 |13085178737| 14010901    |
 |13098336049| 14010901    |
 |13098486609| 14010901    |
 |13097770966| 14010901    |
 |13099744296| 14010901    |
 |13101233891| 14010901    |
 |13100358276| 14010901    |
 +-----------+-------------+

期望结果：

+-----------+-------------+--------------+
 |id         | persian_date| date_miladi  |
 +-----------+-------------+--------------+
 |13085178737| 14010901    |2022-11-22    |
 |13098336049| 14010901    |2022-11-22    |
 |13098486609| 14010901    |2022-11-22    |
 |13097770966| 14010901    |2022-11-22    |
 |13099744296| 14010901    |2022-11-22    |
 |13101233891| 14010901    |2022-11-22    |
 |13100358276| 14010901    |2022-11-22    |
 +-----------+-------------+--------------+

错误原因分析 #

1. 类型不匹配：persian_date列实际是Decimal类型，直接用下标访问x[1]会报错，需先转为字符串；
2. 返回类型错误：尝试1中UDF返回类型定义为long，与日期输出不符；尝试2中仅返回单个日期对象，而非完整的Series；
3. 迭代器使用错误：尝试3中yield单个字符串，而非Spark期望的Series对象，导致空值检查失败。

正确实现方式 #

首先确保安装波斯日期处理库：pip install jdatetime（功能与JalaliDate一致）

方法1：Scalar型Pandas UDF（适合中小数据量） #

import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StringType
import jdatetime

@pandas_udf(StringType())
def persian_to_gregorian(v: pd.Series) -> pd.Series:
    # 将Decimal类型转为字符串，统一处理格式
    str_dates = v.astype(str)
    # 批量拆分年、月、日
    years = str_dates.str.slice(0, 4).astype(int)
    months = str_dates.str.slice(4, 6).astype(int)
    days = str_dates.str.slice(6, 8).astype(int)
    # 批量转换并格式化为YYYY-MM-DD字符串
    return pd.Series(
        [jdatetime.date(y, m, d).togregorian().strftime("%Y-%m-%d") 
         for y, m, d in zip(years, months, days)]
    )

# 应用UDF到DataFrame
df = df.withColumn("date_miladi", persian_to_gregorian(df.persian_date))
df.show()

方法2：Scalar Iter型Pandas UDF（适合大数据量，内存友好） #

import pandas as pd
from typing import Iterator
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StringType
import jdatetime

@pandas_udf(StringType())
def persian_to_gregorian_iter(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for series in iterator:
        str_dates = series.astype(str)
        years = str_dates.str.slice(0, 4).astype(int)
        months = str_dates.str.slice(4, 6).astype(int)
        days = str_dates.str.slice(6, 8).astype(int)
        # 生成转换后的结果Series
        result = pd.Series(
            [jdatetime.date(y, m, d).togregorian().strftime("%Y-%m-%d")
             for y, m, d in zip(years, months, days)]
        )
        yield result

# 应用UDF到DataFrame
df = df.withColumn("date_miladi", persian_to_gregorian_iter(df.persian_date))
df.show()

关键注意事项 #

• 类型统一：先将Decimal转为字符串，避免下标访问错误；
• 批量处理：利用Pandas的矢量化操作（str.slice）替代逐元素map，提升处理效率；
• 返回匹配：确保返回完整的pd.Series，UDF返回类型与输出格式一致（也可改用DateType，需调整返回值为datetime.date对象）；
• 迭代器模式：Scalar Iter型UDF按批次处理数据，减少内存占用，更适合百万级以上数据量。

内容的提问来源于stack exchange，提问作者M_Gh