按字段添加两个结构答案

【问题标题】：Adding up two structs by field按字段添加两个结构
【发布时间】：2022-07-16 02:03:03
【问题描述】：

我有一个如下的 DataFrame 架构（在 PySpark 中）：

root
 |-- ID: string (nullable = true)
 |-- 2022: struct (nullable = true)
 |    |-- val_1: int (nullable = true)
 |    |-- val_2: double (nullable = true)
 |    |-- val_3: double (nullable = true)
 |-- 2021: struct (nullable = true)
 |    |-- val_1: int (nullable = true)
 |    |-- val_2: double (nullable = true)
 |    |-- val_3: double (nullable = true)

我想将2021 和2022 列添加到一列AGG 中，这样它将包含DataFrame 中每一行的2021 和2022 的总和。

我尝试了以下方法：

df.select(
    'ID',
    (F.col("2021") + F.col("2022")).alias("AGG")
).printSchema()

所需的输出架构应该在哪里：

root
 |-- ID: string (nullable = true)
 |-- AGG: struct (nullable = true)
 |    |-- val_1: int (nullable = true)
 |    |-- val_2: double (nullable = true)
 |    |-- val_3: double (nullable = true)

但似乎 spark 只支持添加数字类型。有没有办法隐式添加它们而不直接为结构中的每个字段编写它？

【问题讨论】：

能否添加示例数据框以实现相同的架构和可重现的数据框？

标签： dataframe apache-spark pyspark apache-spark-sql

【解决方案1】：

不，您不能以这种方式添加结构列。您需要通过对嵌套字段求和来创建一个新的结构列。

您可以对内部字段列表使用列表推导来创建一个新列，如下所示：

from pyspark.sql import functions as F


def add_struct(common_fields, s1, s2):
    return F.struct(*[
        (s1[f] + s2[f]).alias(f) for f in common_fields
    ])


# get list of struct fields from schema
fields = df.select("2022.*").columns

df.select(
    'ID',
    add_struct(fields, F.col("2021"), F.col("2022")).alias("AGG")
)

【讨论】：