PySpark中合并两个数组结构体列的技术求助

需求说明 #

将colA和colB两个数组列中同索引位置的结构体元素合并，生成包含合并后结构体的新数组列colA-colB。

解决方案 #

方法1：使用Spark 2.4+高阶函数（推荐）

利用zip_with函数将两个数组的元素按位置配对，再通过struct(a.*, b.*)直接合并对应结构体的所有字段（字段名无冲突时适用）。

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, zip_with, struct

# 初始化Spark会话
spark = SparkSession.builder.appName("MergeArrayStructs").getOrCreate()

# 构造测试数据与Schema
data = [
    (
        [
            {"itemType": "type1", "productCode": "prcd1", "productName": "pen"},
            {"itemType": "type2", "productCode": "Prcd2", "productName": "book"}
        ],
        [
            {"city": "delhi", "country": "india"},
            {"city": "la", "country": "usa"}
        ]
    )
]

schema = """
struct<
    colA: array<struct<itemType:string, productCode:string, productName:string>>,
    colB: array<struct<city:string, country:string>>
>
"""

df = spark.createDataFrame(data, schema=schema)

# 合并对应位置的结构体
merged_df = df.withColumn(
    "colA-colB",
    zip_with(
        col("colA"),
        col("colB"),
        lambda a, b: struct(a.*, b.*)
    )
)

# 查看结果
merged_df.show(truncate=False)

如果两个结构体存在同名字段，需要手动指定字段别名避免冲突：

merged_df = df.withColumn(
    "colA-colB",
    zip_with(
        col("colA"),
        col("colB"),
        lambda a, b: struct(
            a["itemType"], a["productCode"], a["productName"],
            b["city"].alias("b_city"), b["country"].alias("b_country")
        )
    )
)

方法2：自定义UDF（兼容低版本Spark）

如果Spark版本低于2.4，可通过自定义UDF实现合并逻辑：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StructType, StructField, StringType

# 定义合并逻辑：遍历数组，合并对应位置的字典
def merge_struct_pairs(arr_a, arr_b):
    return [dict(**x, **y) for x, y in zip(arr_a, arr_b)]

# 定义返回结果的Schema
result_schema = ArrayType(StructType([
    StructField("itemType", StringType()),
    StructField("productCode", StringType()),
    StructField("productName", StringType()),
    StructField("city", StringType()),
    StructField("country", StringType())
]))

# 注册UDF
merge_udf = udf(merge_struct_pairs, result_schema)

# 生成合并列
merged_df = df.withColumn("colA-colB", merge_udf(col("colA"), col("colB")))
merged_df.show(truncate=False)

注意事项 #

• 确保colA和colB的数组长度一致，否则zip_with会以较短数组的长度为准，超出的元素会被忽略
• 若结构体字段名冲突，必须手动指定别名，否则会抛出字段重复的错误

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