如何高效合并Spark DataFrame中两个不同字段的数组结构体列？

我有一个Spark DataFrame，Schema如下：

StructType(
    [
        StructField('product_id', IntegerType(), True),
        StructField('tenant_id', IntegerType(), True),
        StructField("materials", ArrayType(StructType([
            StructField('id', IntegerType(), True),
            StructField('percentage', FloatType(), True)]
        ))),
        StructField("elastic", ArrayType(StructType([
            StructField('id', IntegerType(), True),
            StructField('name', MapType(StringType(), StringType()), True)
        ])))
    ])

需要将materials和elastic两个数组列按id匹配合并，生成新的materials数组（包含id、percentage、name字段），最终Schema如下：

StructType(
    [
        StructField('product_id', IntegerType(), True),
        StructField('tenant_id', IntegerType(), True),
        StructField("materials", ArrayType(StructType([
            StructField('id', IntegerType(), True),
            StructField('percentage', FloatType(), True),
            StructField('name', MapType(StringType(), StringType()), True)]
        )))
    ])

之前用UDF实现了需求，但性能很差，原UDF代码如下：

@udf(returnType=ArrayType(StructType([
    StructField("id", IntegerType(), False),
    StructField('percentage', FloatType(), True),
    StructField('name', MapType(StringType(), StringType()), True)
])))
def expand_list(materials, elastic):
    final = []
    for k in materials:
        for i in elastic:
            if k.id == i.id:
                final += [{'id': k.id, 'percentage': k.percentage, 'name': i.name}]
    return final

Spark 2.4及以上版本支持数组高阶函数，完全不需要用UDF，能利用Spark的内置优化，性能远优于Python UDF。核心逻辑是遍历materials数组的每个元素，在elastic数组中过滤出id匹配的元素，提取name字段后合并成新结构体。

代码实现：

from pyspark.sql import functions as F

# 生成新的materials数组，合并匹配的name字段
df_result = df.withColumn(
    "materials",
    F.transform(
        "materials",
        lambda material: F.struct(
            material["id"].alias("id"),
            material["percentage"].alias("percentage"),
            # 从elastic中过滤出id匹配的元素，取第一个的name字段
            F.element_at(
                F.filter("elastic", lambda e: e["id"] == material["id"]),
                1
            )["name"].alias("name")
        )
    )
).drop("elastic")

如果需要处理更复杂的匹配逻辑，或者数据量极大，可以先把两个数组展开成行，关联后再重新聚合成数组。这种方式利用Spark的分布式join优化，同样比UDF高效。

代码实现：

from pyspark.sql import functions as F

# 展开materials数组，提取字段
df_materials = df.select(
    "product_id", "tenant_id",
    F.explode("materials").alias("material")
).select(
    "product_id", "tenant_id",
    F.col("material.id").alias("id"),
    F.col("material.percentage").alias("percentage")
)

# 展开elastic数组，提取字段
df_elastic = df.select(
    "product_id", "tenant_id",
    F.explode("elastic").alias("elastic")
).select(
    "product_id", "tenant_id",
    F.col("elastic.id").alias("id"),
    F.col("elastic.name").alias("name")
)

# 按product_id、tenant_id、id关联，保留materials中所有元素（左连接）
df_joined = df_materials.join(
    df_elastic,
    on=["product_id", "tenant_id", "id"],
    how="left"
)

# 重新聚合成目标数组结构
df_result = df_joined.groupBy("product_id", "tenant_id").agg(
    F.collect_list(
        F.struct("id", "percentage", "name")
    ).alias("materials")
)

• 高阶函数方案：代码简洁，无需展开聚合，适合简单的数组内匹配场景，性能最优，Spark会对高阶函数做逻辑优化，避免Python层面的低效循环。
• 展开聚合方案：适合复杂关联逻辑（比如多条件匹配、多表关联），利用Spark的分布式计算能力，比UDF更高效，因为UDF是Python单进程循环，无法利用Spark的集群资源。

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