PySpark：将键值结构体数组转换为单个结构体

需求说明 #

现有一列数据为任意长度的键值结构体数组，需将其转换为预定义的Person结构体列。要求仅使用pyspark.sql.functions内置方法，禁止使用Python UDF，确保能高效处理数十亿行数据。

数据结构定义 #

输入列Schema #

from pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatType, DoubleType, ArrayType

input_schema = StructType([
    StructField("input_column", ArrayType(StructType([
        StructField("key", StringType(), False),
        StructField("value", StructType([
            StructField("string_value", StringType(), True),
            StructField("int_value", IntegerType(), True),
            StructField("float_value", FloatType(), True),
            StructField("double_value", DoubleType(), True)
        ]), True)
    ])), True)
])

目标Person结构体Schema #

person_schema = StructType([
    StructField("first_name", StringType(), True),
    StructField("middle_name", StringType(), True),
    StructField("last_name", StringType(), True),
    StructField("birth_year", IntegerType(), True),
    StructField("ssn", IntegerType(), True),
    StructField("zipcode", IntegerType(), True),
])

输入输出示例 #

输入示例 #

[
    (key: "first_name", value: (string_value: "John")),
    (key: "ssn",        value: (int_value:    123456789)),
    (key: "last_name",  value: (string_value: "Doe")),
]
------------------------------------------------------
[
    (key: "ssn",        value: (int_value:    987654321)),
    (key: "last_name",  value: (string_value: "Jones")),
]
------------------------------------------------------
[
    (key: "zipcode",    value: (int_value:    13579)),
    (key: "first_name", value: (string_value: "Bob")),
    (key: "birth_year", value: (int_value:    1985)),
    (key: "last_name",  value: (string_value: "Smith")),
]

期望输出 #

{
    first_name: "John",
    last_name:  "Doe",
    ssn:        123456789
}
------------------------------------------------------
{
    last_name:  "Jones",
    ssn:        987654321
}
------------------------------------------------------
{
    first_name: "Bob",
    last_name:  "Smith",
    birth_year: 1985,
    zipcode:    13579
}

实现方案 #

核心思路是先将键值数组转换为Map类型，再从Map中提取对应字段组装成目标结构体，全程使用Spark内置矢量化函数保证性能。

完整代码 #

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatType, DoubleType, ArrayType
from pyspark.sql import functions as F

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

# 定义输入Schema
input_schema = StructType([
    StructField("input_column", ArrayType(StructType([
        StructField("key", StringType(), False),
        StructField("value", StructType([
            StructField("string_value", StringType(), True),
            StructField("int_value", IntegerType(), True),
            StructField("float_value", FloatType(), True),
            StructField("double_value", DoubleType(), True)
        ]), True)
    ])), True)
])

# 构造测试数据
test_data = [
    ([
        {"key": "first_name", "value": {"string_value": "John", "int_value": None, "float_value": None, "double_value": None}},
        {"key": "ssn", "value": {"string_value": None, "int_value": 123456789, "float_value": None, "double_value": None}},
        {"key": "last_name", "value": {"string_value": "Doe", "int_value": None, "float_value": None, "double_value": None}}
    ],),
    ([
        {"key": "ssn", "value": {"string_value": None, "int_value": 987654321, "float_value": None, "double_value": None}},
        {"key": "last_name", "value": {"string_value": "Jones", "int_value": None, "float_value": None, "double_value": None}}
    ],),
    ([
        {"key": "zipcode", "value": {"string_value": None, "int_value": 13579, "float_value": None, "double_value": None}},
        {"key": "first_name", "value": {"string_value": "Bob", "int_value": None, "float_value": None, "double_value": None}},
        {"key": "birth_year", "value": {"string_value": None, "int_value": 1985, "float_value": None, "double_value": None}},
        {"key": "last_name", "value": {"string_value": "Smith", "int_value": None, "float_value": None, "double_value": None}}
    ],)
]

df = spark.createDataFrame(test_data, schema=input_schema)

# 执行转换
df_transformed = df.withColumn(
    # 转换数组元素：提取key和对应类型的value
    "transformed_entries",
    F.transform(
        "input_column",
        lambda x: F.struct(
            x.key.alias("key"),
            F.when(x.key.isin("first_name", "middle_name", "last_name"), x.value.string_value)
             .when(x.key.isin("birth_year", "ssn", "zipcode"), x.value.int_value)
             .alias("value")
        )
    )
).withColumn(
    # 将转换后的数组转为Map
    "data_map",
    F.map_from_entries(F.col("transformed_entries"))
).withColumn(
    # 从Map中提取字段组装成Person结构体
    "person",
    F.struct(
        F.col("data_map").getItem("first_name").alias("first_name"),
        F.col("data_map").getItem("middle_name").alias("middle_name"),
        F.col("data_map").getItem("last_name").alias("last_name"),
        F.col("data_map").getItem("birth_year").cast(IntegerType()).alias("birth_year"),
        F.col("data_map").getItem("ssn").cast(IntegerType()).alias("ssn"),
        F.col("data_map").getItem("zipcode").cast(IntegerType()).alias("zipcode")
    )
).drop("transformed_entries", "data_map")

# 查看结果
df_transformed.select("person").show(truncate=False)

性能说明 #

• 全程使用Spark内置的transform、map_from_entries、struct等矢量化函数，避免了Python UDF的序列化/反序列化开销，适合处理大规模数据。
• 逻辑清晰，每一步都是Spark优化后的算子，能充分利用集群资源并行处理。

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