PySpark计算DataFrame数组列分组均值问题求助

我正从Python Pandas转用PySpark，现有一个DataFrame，其中sorted_zipped列是通过arrays_zip生成的结构体数组。需求为：按数组中每个结构体的第一个元素（品牌）分组，计算第二个元素的均值，生成新列且结果按均值降序排列。

尝试用UDF实现时遇到两个问题：

• 在Databricks中无法在Worker节点使用SparkContext
• 改用numpy的UDF触发TypeError: cannot perform reduce with flexible type错误

数据示例 #

+----------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|WeekOfYear|sorted_zipped                                                                                                                                                                                                                                                                                         |
+----------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|13-2023   |[{bmw, 0.99}, {vw, 0.98}, {chevy, 0.97}, {buick, 0.96}]                                                                                                                                                                                                                                               |
|14-2023   |[{chevy, 0.98}, {bmw, 0.98}, {bmw, 0.978}, {bmw, 0.976}, {vw, 0.975}, {bmw, 0.975}, {bmw, 0.97}, {buick, 0.967}, {vw, 0.964}, {vw, 0.96}, {nissan, 0.96}, {chevy, 0.952}, {nissan, 0.95}, {nissan, 0.95}, {lexus, 0.95}, {lexus, 0.94}, {lexus, 0.94}, {nissan, 0.935}, {buick, 0.93}, {chevy, 0.928}]
|15-2023   |[{chevy, 0.992}, {bmw, 0.987}, {nissan, 0.982}, {bmw, 0.982}, {buick, 0.978}, {lexus, 0.976}, {bmw, 0.975}, {bmw, 0.97}, {chevy, 0.967}, {vw, 0.964}, {lexus, 0.961}, {nissan, 0.96}, {vw, 0.952}, {nissan, 0.952}, {vw, 0.952}, {lexus, 0.943}]
+----------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

预期输出（第3行） #

[{chevy, 0.9795}, {bmw, 0.9784}, {buick, 0.978}, {nissan, 0.967}, {lexus, 0.96}, {vw, 0.956}]

尝试的错误代码 #

第一种UDF尝试（SparkContext访问问题） #

def get_avg_1(x):
  rdd = parallelize(x)
  rdd2 = rdd.flatMap(lambda x: [(k, v) for (k, v) in x.items()]).collect()
  grouped_k = rdd2.groupByKey()
  avg_map = grouped_k.mapValues(lambda x: sum(x[1])/len(x[1])).collect()
  return avg_map

第二种numpy UDF尝试（类型错误） #

import numpy as np
import json

schema = ArrayType(StructType([
  StructField("GroupedBrands", StringType(), True),
  StructField("GroupedWeights", FloatType(), True)
]))

array_mean = F.udf(lambda x: (x[0], np.mean(x[1]), schema))
mean_df = sdf.withColumn("mean_value", array_mean("sorted_zipped"))

mean_df.show()

正确解决方案（使用Spark内置函数，避免UDF） #

Spark内置函数在分布式环境下更高效，且不会出现Worker节点权限或类型兼容问题，步骤如下：

from pyspark.sql import functions as F
from pyspark.sql.types import *

# 1. 展开结构体数组为单行记录
exploded_df = sdf.withColumn("struct_element", F.explode("sorted_zipped"))

# 2. 提取品牌和数值字段（注意：如果结构体字段是默认的_1、_2，替换成对应名称）
extracted_df = exploded_df.select(
    "WeekOfYear",
    F.col("struct_element._1").alias("brand"),  # 替换为你的品牌字段名
    F.col("struct_element._2").alias("weight")  # 替换为你的数值字段名
)

# 3. 按周+品牌分组，计算均值
grouped_df = extracted_df.groupBy("WeekOfYear", "brand")\
    .agg(F.round(F.avg("weight"), 4).alias("avg_weight"))

# 4. 按周聚合回结构体数组，并按均值降序排序
result_df = grouped_df.groupBy("WeekOfYear")\
    .agg(F.collect_list(
        F.struct(
            F.col("brand").alias("GroupedBrands"),
            F.col("avg_weight").alias("GroupedWeights")
        )
    ).alias("mean_value"))\
    .withColumn("mean_value", F.sort_array("mean_value", asc=False))

result_df.show(truncate=False)

关键说明 #

• 字段名适配：如果你的结构体字段不是brand和weight，而是arrays_zip生成的默认名称（如_1、_2），需要替换代码中对应的字段名
• 排序逻辑：sort_array函数直接对结构体数组按avg_weight降序排列，无需额外处理
• 性能优势：完全使用Spark内置函数，避免UDF带来的序列化/反序列化开销

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