PySpark高阶函数如何一次对结构体数组中两个字段分别求和？

你可以通过将aggregate的累加器定义为结构体类型，仅遍历数组一次就同时计算两个字段的总和，一步得到最终结果，无需拆分两次聚合也无需生成中间临时列，代码如下：

import pyspark.sql.functions as F

# 一步计算最终结果
df = df.withColumn("output", 
    F.aggregate(
        col = "vals",
        # 初始累加器为结构体：分别存储col2总和、col1总和
        initialValue = F.struct(F.lit(0.0).alias("sum_col2"), F.lit(0).alias("sum_col1")),
        # 遍历每个元素时同步更新两个累加值
        merge = lambda acc, element: F.struct(
            (acc.sum_col2 + element.col2).alias("sum_col2"),
            (acc.sum_col1 + element.col1).alias("sum_col1")
        )
    ).sum_col2 / F.col("output").sum_col1
)

说明 #

• 相较于你原本两次调用aggregate的写法，该方案仅对数组做1次遍历，在数组元素较多的场景下性能优势更明显
• 如果你不需要保留中间求和结果，该写法没有额外的中间列生成，代码更简洁

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