Spark非UDF实现：按组聚合ArrayType列并计算中位数

原始数据与需求 #

给定包含ArrayType列的Spark DataFrame：

simpleData = [("202305","Sales","NY",[1,2,3]),
    ("202306","Sales","NY",[4,2,3]),
    ("202305","Sales","CA",[4,5,3]),
    ("202306","Finance","CA",[4,5,6]),
    ("202305","Finance","NY",[5,6,7]),
    ("202306","Finance","NY",[6,7,8]),
  ]

schema = ["month","department","state","lis"]
df = spark.createDataFrame(data=simpleData, schema = schema)

需求：按department分组，将每组中lis列的所有元素合并后计算中位数，要求不使用UDF（避免性能损耗），最终实现类似目标调用逻辑的效果。

预期输出：

+----------+----------+
|department|lis_median|
+----------+----------+
|     Sales|         3|
|   Finance|         6|
+----------+----------+

无UDF实现方案 #

核心思路是先将ArrayType列拆分为单行元素，再利用Spark内置的percentile_approx函数直接计算分位数（中位数即50%分位），步骤如下：

1. 使用explode函数将每个lis数组的元素拆分为独立行
2. 按目标分组列（department）聚合，调用percentile_approx计算0.5分位值

代码实现 #

from pyspark.sql.functions import explode, percentile_approx

# 拆分数组元素为单行
exploded_df = df.select("department", explode("lis").alias("lis_element"))

# 分组计算中位数
result_df = exploded_df.groupBy("department").agg(
    percentile_approx("lis_element", 0.5).alias("lis_median")
)

result_df.show()

结果验证 #

执行上述代码后，输出与预期完全一致。

方案说明 #

• explode是Spark内置的高效列处理函数，避免了嵌套列表的复杂操作
• percentile_approx是近似分位数计算函数，适合大数据场景，性能远高于自定义UDF
• 该方案完全基于Spark原生API，无需自定义UDF，兼顾了效率与简洁性

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