PySpark SQL如何从列名列表获取greatest()值及对应列名

问题1：原写法是否正确？ #

原写法功能上能跑出你要的final_*结果，但存在两处明显问题：

• 语法错误：你最后一行df = (df.withColumn(...)少了闭合的右括号，直接运行会报错
• 逻辑冗余：你把withColumn放在了内层ls2的循环里，相当于每拼接一个列名就修改一次df，比如处理a的时候，会依次用['ad']、['ad','ae']、['ad','ae','af']三次调用greatest修改df，完全可以等new_lst把所有对应列名都收集完之后，在外层循环只调用一次withColumn即可，运行效率更高。

修正后的最简实现（只生成最大值列的版本）：

from pyspark.sql.functions import greatest

lst1 = ["a", "b"]
lst2 = ["d", "e", "f"]
for ls1 in lst1:
    cols = [ls1 + ls2 for ls2 in lst2]
    df = df.withColumn(f"final_{ls1}", greatest(*cols))

运行后就能得到你给出的期望输出结果。

问题2：如何同步获取最大值对应的列名？ #

可以用struct结构绑定列值和列名，利用struct排序时先按第一个元素排序、再按第二个元素排序的特性，取最大的struct就能同时拿到最大值和对应的列名，示例代码如下：

from pyspark.sql.functions import greatest, lit, col, struct

lst1 = ["a", "b"]
lst2 = ["d", "e", "f"]
for ls1 in lst1:
    # 把每个列的值和列名绑成struct：(值, 列名)
    struct_cols = [struct(col(ls1+ls2), lit(ls1+ls2)) for ls2 in lst2]
    # 取最大的struct
    max_struct = greatest(*struct_cols)
    # 分别提取最大值和对应的列名
    df = df.withColumn(f"final_{ls1}", max_struct["col1"]) \
           .withColumn(f"final_{ls1}_col", max_struct["col2"])

运行后除了你要的final_a、final_b最大值列，还会多出final_a_col、final_b_col两列，分别存对应最大值的列名，比如第一行final_a_col的值就是af。
如果存在多个列值相同且都是最大值的情况，greatest会返回排序靠前的那个列，你也可以根据需求调整逻辑处理并列最大值的场景。

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