PySpark如何获取DataFrame每行最大值对应的列索引

核心实现思路 #

• 先筛选出需要参与最大值比较的列（排除ID列）
• 为每个待比较列绑定对应的索引值（Col1对应1，以此类推）
• 利用PySpark内置数组函数找到每行最大值对应的索引

完整代码示例 #

首先导入依赖的函数：

from pyspark.sql import functions as F

假设你的原始DataFrame变量名为df，后续操作如下：

# 1. 筛选参与最大值比较的列，排除ID列
compare_cols = [col for col in df.columns if col != "ID"]
# 2. 构造(列值, 列索引)的结构体数组，索引从1开始匹配需求
value_index_array = F.array(*[
    F.struct(F.col(c).alias("value"), F.lit(idx + 1).alias("index"))
    for idx, c in enumerate(compare_cols)
])
# 3. 提取最大值对应的索引作为MAX列
df_result = df.withColumn("MAX", F.array_max(value_index_array).index)

补充说明 #

如果存在多个列数值相同且均为最大值的情况，上述代码会返回第一个出现的最大值对应的列索引。如果需要返回最后一个出现的索引，可以自定义UDF处理数组逻辑即可。

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