在 pyspark DataFrame 中添加两个结构字段

Question

我有一个 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 中，这样它将包含每个 2021 和 2022 的总和DataFrame 中的行。

我尝试了以下方法：

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好像只支持添加数值类型。有没有办法隐式添加它们而不直接为结构中的每个字段写入它们？

Answer 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")
)

在 pyspark DataFrame 中添加两个结构字段

Adding two struct fields in pyspark DataFrame

dataframe

apache-spark

pyspark

apache-spark-sql