PySpark按行值降序生成列名数组新列的Scala转Python问题

直接上对应Python实现代码：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, array, col
from pyspark.sql.types import ArrayType, StringType

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

# 创建示例DataFrame
df = spark.createDataFrame([(1,4,3), (4,1,3)], ["a", "b", "c"])

# 获取所有列名
names = df.columns

# 定义UDF：接收值的序列，和列名配对后按值降序排序，提取列名
def sort_names(values):
    # 配对值和列名，按值降序排序，然后取列名组成数组
    return [col_name for _, col_name in sorted(zip(values, names), key=lambda x: x[0], reverse=True)]

# 注册UDF，指定返回类型为字符串数组
sort_names_udf = udf(sort_names, ArrayType(StringType()))

# 添加新列
result_df = df.withColumn("newcol", sort_names_udf(array(*[col(name) for name in names])))

# 展示结果
result_df.show(truncate=False)

代码说明 #

1. 列名获取：Python中用df.columns替代Scala的df.schema.fieldNames，效果完全一致。
2. UDF定义：Python UDF无需像Scala那样显式声明输入类型，直接定义函数即可，注册时指定返回类型为ArrayType(StringType)保证类型匹配。
3. 排序逻辑：原Scala代码是升序排序，按需求调整为降序，通过sorted函数的reverse=True参数实现，也可以用key=lambda x: -x[0]达到同样效果。
4. 数组构造：Python中array(*[col(name) for name in names])对应Scala的array(names.map(col): _*)，将所有列打包成数组传递给UDF。

执行后输出结果与示例一致：

+---+---+---+---------+
|a  |b  |c  |newcol   |
+---+---+---+---------+
|1  |4  |3  |[b, c, a]|
|4  |1  |3  |[a, c, b]|
+---+---+---+---------+

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