Spark DataFrame分组生成字典列表列并通过UDF批量处理

import sys
import boto3
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from awsglue.dynamicframe import DynamicFrame
from pyspark.sql.functions import regexp_replace, col, expr, collect_list, struct, udf
from pyspark.sql.types import ArrayType, MapType, StringType, IntegerType

args = getResolvedOptions(sys.argv, ['JOB_NAME'])

sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session

job = Job(glueContext)
job.init(args["JOB_NAME"], args)

# 初始化数据
d = [
 {"account_number": 1, "v1": 100830, "v2": 1000},
 {"account_number": 2, "v1": 2000, "v2": None},
 {"account_number": 2, "v1": 4545, "v2": 55},
 {"account_number": 3, "v1": 4545, "v2": 55},
 {"account_number": 3, "v1": 555, "v2": 55}
]

df = spark.createDataFrame(d)

# 1. 新增列：按account_number分组，存储同一账号下所有行的字典列表
# 先将每行转成struct，再收集为列表，后续可转为字典格式
account_grouped = df.groupBy("account_number")\
    .agg(collect_list(struct(*df.columns)).alias("account_rows"))

# 定义业务UDF，示例逻辑为计算账号下v1的总和，可按需修改
def process_account_rows(rows):
    # 将struct对象转为字典列表
    row_dicts = [row.asDict() for row in rows]
    # 自定义处理逻辑：计算v1总和（忽略空值）
    total_v1 = sum(row["v1"] for row in row_dicts if row["v1"] is not None)
    return {"total_v1": total_v1}

# 注册UDF，需根据实际返回结果调整类型定义
process_udf = udf(process_account_rows, MapType(StringType(), IntegerType()))

# 2. 每个账号仅调用一次UDF，得到处理结果
account_processed = account_grouped.withColumn("udf_result", process_udf(col("account_rows")))

# 将处理结果关联回原DataFrame，更新对应账号的所有行
final_df = df.join(account_processed, on="account_number", how="left")

# 查看最终结果
final_df.show(truncate=False)

job.commit()

• 生成账号行列表：通过groupBy+collect_list(struct(*df.columns))，把同一账号的所有行聚合为struct列表，后续可直接转为字典列表供UDF处理。
• 减少UDF调用次数：仅在分组聚合后的小数据集上调用UDF，每个账号仅执行一次计算，避免重复调用带来的性能损耗，再通过join把结果关联回原表的所有对应行。
• UDF自定义扩展：process_account_rows函数是示例逻辑，你可以根据业务需求替换为任意处理逻辑，同时注意同步调整UDF的返回类型定义。

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