添加新列,其中包含按行值pyspark排序的列名数组

h9vpoimq  于 2023-01-12  发布在  Spark
关注(0)|答案(1)|浏览(126)

我是Pyspark的新手,我发现很久以前就有人问过这个问题了
给定一个包含几个列的dataFrame,我尝试创建一个新列,其中包含一个数组,该数组中的列名称根据这些列的行值按降序排序。

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

这是他们针对Scala的解决方案

val df = spark.createDataFrame(Seq((1,4,3), (4,1,3))).toDF("a", "b", "c")

val names=df.schema.fieldNames
val sortNames = udf((v: Seq[Int]) => {v.zip(names).sortBy(_._1).map(_._2)})
df.withColumn("newcol", sortNames(array(names.map(col): _*))).show

我在把这段代码翻译成Python时遇到了麻烦

ffx8fchx

ffx8fchx1#

这一个应该做到这一点:

from pyspark.sql.functions import udf, array

df = spark.createDataFrame([(1, 4, 3), (4, 1, 3)], "a: int, b: int, c: int")
colNames = df.columns

@udf
def orderByColValue(colValues: list):
    return [x[1] for x in sorted(zip(colValues, colNames), key=lambda x: x[0], reverse=True)]

df.withColumn("ordered_cols", orderByColValue(array(colNames))).show(truncate=False)

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

相关问题