PySpark如何向ArrayType列中条件添加元素？

问题背景 #

给定以下PySpark数据结构：

from pyspark.sql import Row
from pyspark.sql.types import StructType, StructField, StringType, ArrayType

ItemStruct = StructType([StructField("BomId", StringType()), StructField("price", StringType())])
BomStruct = StructType([StructField("OrderId", StringType()), StructField("items", ArrayType(ItemStruct))])
sampledata_sof = [Row("123-A", [Row("Bom-11", "120"), Row("Bom-12", "140")]), Row("100-A", [Row("Bom-23", "170"), Row("Bom-24", "190")])]

dfSampleBom = spark.createDataFrame(spark.sparkContext.parallelize(sampledata_sof), BomStruct)
dfSampleBom.printSchema() 
dfSampleBom.show()

需求：当items数组中存在Bom-11时，向该数组追加元素Row("Bom-99", "99")（注意price字段为StringType，需传入字符串值）。

曾尝试使用df.rdd.map(lambda x: generateItems(x))处理，触发报错：

pyspark.errors.exceptions.base.PySparkRuntimeError: [CONTEXT_ONLY_VALID_ON_DRIVER] It appears that you are attempting to reference SparkContext from a broadcast variable, action, or transformation. SparkContext can only be used on the driver, not in code that it run on workers. For more information, see SPARK-5063.

需要Spark原生支持的高效分布式处理方案，不确定UDF是否可行。

报错原因 #

报错根源是在RDD的map转换中（或generateItems函数内部）直接引用了SparkContext对象。Spark的分布式转换逻辑在Worker节点执行，而SparkContext仅能在Driver节点访问，因此触发权限限制报错。

解决方案 #

方案1：Spark原生函数（推荐，高效分布式） #

使用Spark内置函数实现条件判断与数组追加，无需自定义UDF，性能最优：

from pyspark.sql import functions as F

# 处理逻辑：检查items中是否包含Bom-11，满足则追加元素
df_result = dfSampleBom.withColumn(
    "items",
    F.when(
        # 提取items中的所有BomId，判断是否包含目标值
        F.array_contains(F.transform("items", lambda item: item.BomId), "Bom-11"),
        # 合并原数组与新构造的元素数组
        F.concat(
            F.col("items"),
            F.array(F.struct(F.lit("Bom-99").alias("BomId"), F.lit("99").alias("price")))
        )
    ).otherwise(F.col("items"))  # 不满足条件则保留原数组
)

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

预期输出：

+-------+----------------------------------------------------+
|OrderId|items                                               |
+-------+----------------------------------------------------+
|123-A  |[{Bom-11, 120}, {Bom-12, 140}, {Bom-99, 99}]        |
|100-A  |[{Bom-23, 170}, {Bom-24, 190}]                      |
+-------+----------------------------------------------------+

方案2：自定义UDF（适合复杂逻辑） #

若需要更灵活的业务逻辑，可使用UDF，但需注意不要在UDF内部引用SparkContext或Driver端专属对象：

# 定义UDF逻辑
def add_bom_99(items):
    # 检查是否存在Bom-11
    has_bom11 = any(item.BomId == "Bom-11" for item in items)
    if has_bom11:
        # 追加新元素（注意price为字符串类型）
        items.append(Row("Bom-99", "99"))
    return items

# 注册UDF，指定返回类型
add_bom_udf = F.udf(add_bom_99, ArrayType(ItemStruct))

# 应用UDF处理
df_result_udf = dfSampleBom.withColumn("items", add_bom_udf("items"))

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

注意：UDF基于Python代码执行，需要序列化/反序列化数据，性能略低于原生函数方案，仅在原生函数无法覆盖复杂逻辑时使用。

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