pyspark agg与groupby不能使用pandas_udf

如何解决pyspark agg与groupby不能使用pandas_udf

当我使用min，sum，max，count，mean或std之一时，pyspark groupby（）。agg（）工作正常。但是如果我 提供用户定义的中值_udf，range_udf或quantum_udf中的任何一个函数，则pyspark groupby（）。agg（）无法正常工作（AttributeError：“ function”对象没有属性“ _get_object_id”）。

import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf,PandasUDFType
from pyspark.sql.types import DateType

spark = SparkSession.builder.getOrCreate()


@pandas_udf("double",PandasUDFType.GROUPED_AGG)
def median_udf(x):
    return np.median(x)


@pandas_udf("double",PandasUDFType.GROUPED_AGG)
def range_udf(x):
    return float(x.max() - x.min())


@pandas_udf("double",PandasUDFType.GROUPED_AGG)
def quantile_udf(x):
    return x.quantile(.50)


df = spark.createDataFrame(
        [
            ('2017-01-01 1:00:00',1000,'IC',280,11),('2017-03-15 8:15:00','ICE',2),('2019-03-27 5:00:00',1002,260,3),('2019-04-01 11:00:00',220,9),('2019-08-07 1:00:00',270,5),('2020-02-22 11:00:00',55),8),('2020-04-08 9:00:00',4),],['Date','TrainID','Traintype','Max_Speed','Delay']
    )

df = df.withColumn("Date",df['Date'].cast(DateType()))

columns = ['Date','Traintype']

# agg_func working - min,sum,max,count,mean,std
# this is working
agg_func = {'Max_Speed': 'min','Delay': 'max'}

# this is not working
# agg_func = {'Max_Speed': 'min','Delay': range_udf}

df = df.groupby(columns).agg(agg_func)
print('Pyspark group by result')
df.show()

我的pandas_udf函数是否存在任何错误？我必须在哪里进行更改才能使其正常工作。