PySpark如何对嵌套Struct字段colB内的所有字段求和？

问题描述 #

现有如下DataFrame Schema：

root
 |-- DataColumn1: struct (nullable = true)
 |    |-- colA: double (nullable = true)
 |    |-- colB: struct (nullable = true)
 |    |    |-- fieldA: double (nullable = true)
 |    |    |-- fieldB: double (nullable = true)
 |    |    |-- fieldC: double (nullable = true)
 |    |-- colC: long (nullable = true)
 |    |-- colD: string (nullable = true)
 |-- DataColumn2: string (nullable = true)
 |-- DataColumn3: string (nullable = true)

需要新增列DataColumn4，值为DataColumn1.colB这个Struct下所有字段的和，且无需逐个指定字段（适配colB包含N个字段的场景），要求不使用Pandas UDF。

示例输入数据：

DataColumn1                 DataColumn2             DataColumn3
(1, (1, 2, 3), 4, 5)           XXX                      YYY
(1, (2, 3, 3), 8, 9)           XYZ                      XYX

预期输出结果：

DataColumn1             DataColumn2             DataColumn3   DataColumn4
(1, (1, 2, 3), 4, 5)       XXX                      YYY         6 (1+2+3=6)
(1, (2, 3, 3), 8, 9)       XYZ                      XYX          8 (2+3+3=8)

解决方案 #

步骤说明 #

1. 提取Struct字段名：从DataFrame的Schema中获取DataColumn1.colB的所有子字段名称。
2. 动态构建求和表达式：遍历字段名，生成每个字段的引用表达式，统一转为DoubleType后求和。
3. 添加新列：将求和表达式作为新列DataColumn4添加到DataFrame中。

代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

# 初始化SparkSession
spark = SparkSession.builder.appName("StructSum").getOrCreate()

# 构建示例数据
data = [
    ((1.0, (1.0, 2.0, 3.0), 4, "5"), "XXX", "YYY"),
    ((1.0, (2.0, 3.0, 3.0), 8, "9"), "XYZ", "XYX")
]

schema = "DataColumn1 struct<colA:double, colB:struct<fieldA:double,fieldB:double,fieldC:double>, colC:long, colD:string>, DataColumn2 string, DataColumn3 string"
df = spark.createDataFrame(data, schema)

# 获取colB的所有字段名
colb_fields = [field.name for field in df.schema["DataColumn1"].dataType["colB"].fields]

# 构建求和表达式：将每个字段转为double后相加
sum_expr = "+".join([f"cast(DataColumn1.colB.{field} as double)" for field in colb_fields])

# 添加新列DataColumn4
result_df = df.withColumn("DataColumn4", expr(sum_expr))

# 展示结果
result_df.show(truncate=False)

代码解释 #

• 提取字段名：通过df.schema["DataColumn1"].dataType["colB"].fields获取colB的所有子字段对象，再提取名称列表，适配任意数量的子字段。
• 动态表达式：用字符串拼接生成类型转换后的求和表达式，确保不同数值类型都能参与计算。
• 执行求和：使用expr()函数解析拼接的表达式，快速生成目标新列。

扩展说明 #

如果colB的子字段可能存在null值，可在表达式中添加coalesce处理，避免求和结果为null：

sum_expr = "+".join([f"coalesce(cast(DataColumn1.colB.{field} as double), 0)" for field in colb_fields])

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