PySpark分组获取每组TopN数据报错的解决方法咨询

嘿，我来帮你搞定这个问题！你遇到的AttributeError: 'function' object has no attribute '_get_object_id'报错，根源是UDF的注册和使用逻辑出了问题，而且其实Spark有更高效的内置方案来实现分组取TopN，不用硬写UDF，咱们一步步来捋清楚：

先说说你代码里的问题 #

你这段UDF注册的代码有两个小错误：

h_f = udf(f, IntegerType())
h_f = spark.udf.register("h_f", h_f)

1. UDF返回类型错误：你的函数f返回的是一个排序后的列表（比如[12,10]），但你给UDF指定的返回类型是IntegerType()，这明显不匹配，应该用ArrayType(IntegerType())。
2. 注册后覆盖了UDF对象：spark.udf.register返回的是一个普通Python函数，不是Spark的UDF对象，你把它重新赋值给h_f后，后面在agg里调用的就不是UDF了，自然会报错。

方法一：修复UDF的写法（如果一定要用UDF） #

如果坚持要用自定义UDF实现，调整下代码就行：

from pyspark.sql import SparkConf, SparkContext, SparkSession
from pyspark.sql.functions import udf, expr
from pyspark.sql.types import ArrayType, IntegerType

# 初始化Spark环境
conf = SparkConf().setMaster("local[*]").setAppName("test")
sc = SparkContext(conf=conf).getOrCreate()
spark = SparkSession(sparkContext=sc).builder.getOrCreate()

# 构建测试数据
rdd = sc.parallelize([(1, 10), (3, 11), (1, 8), (1, 12), (3, 7), (3, 9)])
data = spark.createDataFrame(rdd, ['x', 'y'])
data.show()

# 自定义取Top2的函数
def get_top2(y_list):
    return sorted(y_list, reverse=True)[:2]

# 定义UDF，注意返回类型是数组
top2_udf = udf(get_top2, ArrayType(IntegerType()))
# 注册UDF（可选，注册后可以用SQL风格的名称调用）
spark.udf.register("get_top2", top2_udf)

# 方式1：直接用UDF对象调用
data.groupBy('x').agg(top2_udf(data['y']).alias('top_2_y')).show()

# 方式2：用注册后的函数名，需要用expr包裹
data.groupBy('x').agg(expr("get_top2(y) as top_2_y")).show()

方法二：用Spark窗口函数（更推荐！性能拉满） #

其实Spark内置的窗口函数是实现分组取TopN的最优解，比UDF高效得多（因为Spark能对窗口函数做优化，UDF是黑盒没法优化），代码也更简洁：

from pyspark.sql import Window
from pyspark.sql.functions import row_number, col, collect_list

# 定义窗口规则：按x分组，按y降序排序
window_spec = Window.partitionBy('x').orderBy(col('y').desc())

# 给每组的每条数据添加行号，然后筛选行号<=2的记录
top2_data = data.withColumn('row_num', row_number().over(window_spec)) \
                .filter(col('row_num') <= 2) \
                .drop('row_num')

top2_data.show()

# 如果想把每组的Top2合并成一个列表，再加一步分组聚合
top2_list = top2_data.groupBy('x').agg(collect_list('y').alias('top_2_y'))
top2_list.show()

运行这段代码，你就能得到每组x对应的Top2 y值，不管是展开的记录还是合并成列表的形式都能实现。

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