带有参数的pyspark udf

Question

需要将一个pyspark数据帧列checkin_time从milisec传输到时区调整时间戳，时区信息在另一列tz_info中。

尝试以下：

def tz_adjust(x,tz_info):
    if tz_info:
        y = col(x)+ col(tz_info) 
        return from_unixtime(col(y)/1000)
    else:
        return from_unixtime(col(x)/1000)
    
def udf_tz_adjust(tz_info):
    return udf(lambda l: tz_adjust(l, tz_info))     

同时使用这个udf到列

df.withColumn('checkin_time',udf_tz_adjust('time_zone')(col('checkin_time')))

got some error:
AttributeError: 'NoneType' object has no attribute '_jvm'

知道将第二列作为参数传递给 udf 吗？ 谢谢。

Answer 1

恕我直言，您正在做的是 UDF 和偏函数的组合，这可能会变得很棘手。我认为您根本不需要将 UDF 用于您的应用程序目的。您可以执行以下操作

#not tested

from pyspark.sql.functions import *

df.withColumn('checkin_time', when(col("tz_info").isNotNull(), (from_unixtime(col('checkin_time')) + F.col("tz_info"))/1000).otherwise(from_unixtime(col("checkin_time"))/1000))

UDF 有其自己的serde 效率低下，这在与 python 一起使用时更糟糕，因为它将 scala 数据类型转换为 python 数据类型会产生额外的开销。