如何更新Spark DataFrame数组列中的指定元素？

要解决这个数组列元素替换的问题，我们可以利用PySpark的内置函数高效实现，无需编写复杂的UDF。下面分两种常见场景给出解决方案：

场景1：全局替换所有行数组中的指定值 #

如果你想把所有行的colleagues数组里的guy2都替换成guy10，可以使用Spark 3.0+提供的transform函数，它能遍历数组的每个元素并应用转换逻辑：

步骤1：构建示例DataFrame #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, transform

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

# 创建你提供的示例数据
data = [
    (["guy1", "guy2", "guy3"], "Thisguy"),
    (["guy4", "guy5", "guy6"], "Thatguy"),
    (["guy7", "guy8", "guy9"], "Someguy")
]

df = spark.createDataFrame(data, ["colleagues", "name"])
df.show(truncate=False)

步骤2：执行全局替换 #

df_updated = df.withColumn(
    "colleagues",
    transform(
        col("colleagues"),
        lambda element: when(element == "guy2", "guy10").otherwise(element)
    )
)

df_updated.show(truncate=False)

执行后输出结果：

+---------------------+-------+
|colleagues           |name   |
+---------------------+-------+
|[guy1, guy10, guy3]  |Thisguy|
|[guy4, guy5, guy6]   |Thatguy|
|[guy7, guy8, guy9]   |Someguy|
+---------------------+-------+

场景2：仅替换特定行中的数组元素 #

如果你的需求是只替换第一行（name为Thisguy）里的guy2，可以结合when条件判断行，再对目标行的数组执行替换：

df_updated = df.withColumn(
    "colleagues",
    when(
        col("name") == "Thisguy",  # 定位目标行
        transform(
            col("colleagues"),
            lambda element: when(element == "guy2", "guy10").otherwise(element)
        )
    ).otherwise(col("colleagues"))  # 其他行保持原数组不变
)

df_updated.show(truncate=False)

这个逻辑会精准只修改name为Thisguy的行，其他行的colleagues数组完全不受影响。

补充说明 #

• transform函数是Spark 3.0及以上版本的特性，性能比自定义UDF好很多，优先推荐使用。
• 如果你的Spark版本低于3.0，也可以用自定义UDF实现，但需要注意UDF的序列化开销，示例如下（不推荐，仅作兼容参考）：

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

def replace_element(arr, target, replacement):
    return [replacement if x == target else x for x in arr]

replace_udf = udf(replace_element, ArrayType(StringType()))

df_updated = df.withColumn(
    "colleagues",
    when(
        col("name") == "Thisguy",
        replace_udf(col("colleagues"), "guy2", "guy10")
    ).otherwise(col("colleagues"))
)

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