无需SQL表达式，基于ID从数组取值添加DataFrame列

我有一个DataFrame df，数据结构如下：

{
    "data":[
        {
            "id":"a",
            "val":1
        },
        {
            "id":"b",
            "val":2
        }
    ]
}

需要添加新列test，使其值为数组中id等于'b'的项的val（即2）。目前通过SQL表达式实现了需求：

import pyspark.sql.functions as F

(
    df
    .withColumn(
        "test",
        F.expr("filter(data,x->x.id=='b')")[0]["val"]
    )
    .show()
)

执行结果符合预期：

+----------------+----+
| data           |test|
+----------------+----+
|[{a, 1}, {b, 2}]| 2  |
+----------------+----+

想知道有没有不依赖SQL、更“原生”的PySpark实现方式？

方法1：使用array_filter（PySpark 3.1+推荐） #

PySpark 3.1及以上版本提供了原生的array_filter函数，无需编写SQL表达式就能完成数组过滤：

import pyspark.sql.functions as F

(
    df
    .withColumn(
        "filtered_data",
        F.array_filter(F.col("data"), lambda x: x["id"] == "b")
    )
    .withColumn("test", F.col("filtered_data")[0]["val"])
    .drop("filtered_data")  # 可选：清理中间临时列
    .show()
)

array_filter接收数组列和过滤逻辑的lambda表达式，直接筛选出id为'b'的元素，之后取第一个元素的val字段即可。

方法2：兼容低版本的transform+aggregate组合 #

如果你的PySpark版本低于3.1，可以用transform标记目标元素，再通过aggregate提取对应值：

import pyspark.sql.functions as F

(
    df
    .withColumn(
        "test",
        F.aggregate(
            F.transform(F.col("data"), lambda x: F.when(x["id"] == "b", x["val"]).otherwise(None)),
            F.lit(None),
            lambda acc, x: F.coalesce(x, acc)
        )
    )
    .show()
)

transform会将符合条件的元素替换为对应的val值，不符合的设为None；aggregate则通过coalesce提取第一个非空值（即目标的2）。

方法3：固定索引直接取值（仅适用于元素顺序固定场景） #

如果能确保id为'b'的元素在数组中的位置固定（比如总是第2个元素），可以直接通过索引取值：

import pyspark.sql.functions as F

(
    df
    .withColumn("test", F.col("data")[1]["val"])
    .show()
)

注意：这种方法依赖数组元素的顺序，一旦数据顺序变化会导致结果错误，不推荐用于动态数据场景。

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