如何在PySpark DataFrame中获取次高值对应的列名

原始数据 #

需求 #

需要新增max_value_2（第二高消费金额）和MAX_2（对应第二高消费的类别）两列，计算次高值时必须排除MAX列中已经标记的最高消费类别列。

现有代码（获取最大值列名） #

cond = "psf.when" + ".when".join(["(psf.col('" + c + "') == psf.col('max_value'), psf.lit('" + c + "'))" for c in df.columns]) 

df = df.withColumn("max_value", psf.greatest(*[c for c in columns])) \
       .withColumn("MAX", when(cond, 1).otherwise(0))

期望输出 #

解决方案 #

下面是修改后的代码，实现排除最高值列后计算次高值及对应类别：

from pyspark.sql import functions as psf

# 先筛选出所有消费相关的列（排除id、max_value、MAX这些非消费列）
trx_columns = [col for col in df.columns if col not in ["id", "max_value", "MAX"]]

# 计算第二高消费金额：排除MAX标记的类别列，取剩余列的最大值
df = df.withColumn(
    "max_value_2",
    psf.greatest(
        *[psf.when(psf.col("MAX") != psf.lit(col), psf.col(col)) for col in trx_columns]
    )
)

# 构建匹配第二高类别的条件逻辑
max2_cond = None
for col in trx_columns:
    # 先判断当前列不是最高类别，再检查该列值是否等于次高值
    current_cond = psf.when(
        (psf.col("MAX") != psf.lit(col)) & (psf.col(col) == psf.col("max_value_2")),
        psf.lit(col)
    )
    if max2_cond is None:
        max2_cond = current_cond
    else:
        max2_cond = max2_cond.when(current_cond._expr, current_cond._value)

# 添加MAX_2列
df = df.withColumn("MAX_2", max2_cond)

# 查看最终结果
df.show()

代码解释 #

1. 筛选消费列：明确区分消费数据列和其他辅助列，避免干扰计算。
2. 计算次高金额：用when语句跳过MAX列标记的最高类别列，只保留其他列的值，再通过greatest取这些值的最大值，得到次高金额。
3. 匹配次高类别：遍历每个消费列，判断该列不是最高类别且值等于次高金额时，返回列名，最终组合成MAX_2列的逻辑。

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