PySpark：从数组列取最大值并获取对应索引的另一列值

给定如下PySpark DataFrame：

|   id      |   label   |  md  |
+-----------+-----------+------+
|[a, b, c]  | [1, 4, 2] |  3   |
|[b, d]     | [7, 2]    |  1   |
|[a, c]     | [1, 2]    |  8   |

需求是从label数组列中获取最大值，并提取id数组列中对应索引位置的元素，最终得到如下结果：

| id |label|  md  |
+----+-----+------+
| b  |  4  |  3   |
| b  |  7  |  1   |
| c  |  2  |  8   |

实现代码 #

使用PySpark高阶函数可高效完成需求，无需展开数组再分组：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 初始化SparkSession
spark = SparkSession.builder.appName("array_max_match").getOrCreate()

# 创建示例DataFrame
data = [
    (["a", "b", "c"], [1, 4, 2], 3),
    (["b", "d"], [7, 2], 1),
    (["a", "c"], [1, 2], 8)
]
df = spark.createDataFrame(data, ["id", "label", "md"])

# 核心处理逻辑
result_df = df.withColumn(
    "max_pair",
    # 将id和label按索引配对成结构体数组，再取label最大的结构体
    F.array_max(
        F.zip_with(
            F.col("id"),
            F.col("label"),
            lambda x, y: F.struct(y.alias("label"), x.alias("id"))
        )
    )
).select(
    # 从最大结构体中提取id、label，保留md列
    F.col("max_pair.id").alias("id"),
    F.col("max_pair.label").alias("label"),
    F.col("md")
)

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

逻辑说明 #

1. zip_with配对数组：将id和label数组按相同索引位置配对，生成包含label和id的结构体数组，比如第一行生成[{"label":1,"id":"a"}, {"label":4,"id":"b"}, {"label":2,"id":"c"}]。
2. array_max取最大值结构体：array_max会根据结构体的字段顺序比较大小，这里把label放在结构体的第一个位置，因此会直接筛选出label值最大的那个结构体。
3. 提取目标字段：从筛选出的最大结构体中提取id和label，同时保留原有的md字段，得到最终结果。

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