PySpark多参数用户定义函数UDF返回NULL值解决方案

问题场景 #

尝试将Python日期金额计算逻辑转换为PySpark UDF，读取制表符分隔的loan.txt文件生成新列NewLoanAmount时，所有计算结果均返回NULL。

问题复现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf,col,array
from pyspark.sql.types import StringType,IntegerType,DecimalType
from datetime import date

def calculateAmount(loandate,loanamount):
    y,m,d = loandate.split('-')[0],loandate.split('-')[1],loandate.split('-')[2]
    ld = date(int(y),int(m),int(d))
    if (date(2010,1,1) <= ld <= date(2015,12,31)):
        fine = 10
    elif (date(2016,1,1) <=ld <= date.today()):
        fine = 5
    return ((100+fine)*int(loanamount))/100

spark = SparkSession.builder.appName("User Defined Functions").getOrCreate()
df = spark.read.options(delimiter = "\t",header = True).csv("../input/applicationloan/loan.txt")

calAmount = udf(lambda interest,amount : calculateAmount(interest,amount),DecimalType())

df = df.withColumn("NewLoanAmount",calAmount(col("loandate"),col("loanamount")))
df.show()

问题现象 #

• 代码运行输出结果：
  
• 源文件loan.txt内容（制表符分隔）：
  

问题根因 #

返回NULL和多参数UDF的传参逻辑无关，核心原因有3个：

• 返回值类型不匹配：Spark读取CSV默认所有列为字符串类型，UDF声明返回DecimalType，但Python函数实际返回原生float/int类型，类型无法自动适配时Spark会静默返回NULL，不会抛出显性报错。
• 逻辑分支缺失：如果贷款日期不在判断的两个时间区间内，fine变量未定义，函数运行会触发NameError，Spark执行UDF时捕获到运行时异常默认返回NULL。
• 日期解析鲁棒性差：直接用字符串切分解析日期，只要日期格式和预期不一致就会触发异常返回NULL。

多参数UDF无特殊语法规则：定义Python函数时声明几个入参，调用UDF时按顺序传入对应数量的DataFrame列即可，原代码中lambda传参的写法本身是正确的。

修复方案 #

方案1：修复原有UDF实现 #

针对问题点逐一修复即可正常运行：

1. 给UDF指定明确的Decimal精度，返回值做类型适配
2. 补充时间区间的兜底分支，避免变量未定义
3. 增加异常捕获，避免单条数据格式问题导致全量返回NULL

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import DecimalType
from datetime import date
from decimal import Decimal

def calculateAmount(loandate_str, loanamount_str):
    try:
        y, m, d = loandate_str.split('-')
        ld = date(int(y), int(m), int(d))
        loanamount = int(loanamount_str)
        if date(2010,1,1) <= ld <= date(2015,12,31):
            fine = 10
        elif date(2016,1,1) <= ld <= date.today():
            fine = 5
        else:
            # 非目标时间区间的兜底规则，可根据业务调整
            fine = 0
        # 显式转换为Decimal类型，匹配UDF声明的返回类型
        return Decimal(str(((100 + fine) * loanamount) / 100))
    except:
        return None

spark = SparkSession.builder.appName("User Defined Functions").getOrCreate()
df = spark.read.options(delimiter="\t", header=True).csv("../input/applicationloan/loan.txt")

# 多参数UDF按顺序传入列即可，无需特殊包装
calAmount = udf(calculateAmount, DecimalType(18,2))
df = df.withColumn("NewLoanAmount", calAmount(col("loandate"), col("loanamount")))
df.show()

方案2：原生Spark函数实现（推荐，性能更高） #

Python UDF需要在JVM和Python进程之间做数据序列化，性能远低于Spark原生函数，这类简单计算完全可以不用UDF实现：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, to_date, when, current_date, lit
from pyspark.sql.types import DecimalType

spark = SparkSession.builder.appName("Native Spark Implementation").getOrCreate()
df = spark.read.options(delimiter="\t", header=True).csv("../input/applicationloan/loan.txt")

df = df.withColumn(
    "parsed_loandate",
    to_date(col("loandate"), "yyyy-MM-dd")
).withColumn(
    "parsed_loanamount",
    col("loanamount").cast(IntegerType())
).withColumn(
    "NewLoanAmount",
    when(
        (col("parsed_loandate") >= lit(date(2010,1,1))) & (col("parsed_loandate") <= lit(date(2015,12,31))),
        (lit(110) * col("parsed_loanamount") / lit(100)).cast(DecimalType(18,2))
    ).when(
        (col("parsed_loandate") >= lit(date(2016,1,1))) & (col("parsed_loandate") <= current_date()),
        (lit(105) * col("parsed_loanamount") / lit(100)).cast(DecimalType(18,2))
    ).otherwise(lit(None).cast(DecimalType(18,2)))
).drop("parsed_loandate", "parsed_loanamount")

df.show()

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