Spark DataFrame：根据列值获取按行排序的列名答案

【问题标题】：Spark DataFrame: get row wise sorted column names based on column valuesSpark DataFrame：根据列值获取按行排序的列名
【发布时间】：2019-07-02 23:44:45
【问题描述】：

对于以下数据框中的每一行，我想根据降序列条目查找列名（作为数组或元组或其他内容）。所以，对于数据框

+---+---+---+---+---+
| ID|key|  a|  b|  c|
+---+---+---+---+---+
|  0|  1|  5|  2|  1|
|  1|  1|  3|  4|  5|
+---+---+---+---+---+

我要找

+---+---+---+---+---+------------------+
| ID|key|  a|  b|  c|descending_columns|
+---+---+---+---+---+------------------+
|  0|  1|  5|  2|  1|           [a,b,c]|
|  1|  1|  3|  4|  5|           [c,b,a]|
+---+---+---+---+---+------------------+

理想情况下，一般来说，我希望能够遍历预先指定的列并根据这些列条目应用函数。这可能看起来像：

import pyspark.sql.functions as f

name_cols = ["a","b","c"]

for col in name_cols: 
    values_ls.append = []
    ...schema specification....
    values_ls.append(f.col(col) ...get column value... )

df1 = df.withColumn("descending_columns", values_ls)

这个问题相当简单，但在 pyspark 中有效实施似乎相当具有挑战性。

我使用的是 pyspark 2.3.3 版。

【问题讨论】：

GCP 上的 Spark 版本 2.3.3

标签： apache-spark pyspark

【解决方案1】：

对于小于 2.4 的 Spark 版本，您可以在没有 udf 的情况下使用 sort_array 和 struct 来实现此目的。

首先获取要排序的列列表

cols_to_sort = df.columns[2:]
print(cols_to_sort)
#['a', 'b', 'c']

现在构建一个包含两个元素的结构 - "value" 和 "key"。 "key" 是列名，"value" 是列值。如果您确保"value" 在struct 中排在第一位，则可以使用sort_array 以您想要的方式对该结构数组进行排序。

数组排序后，您只需对其进行迭代并提取包含列名的"key"部分。

from pyspark.sql.functions import array, col, lit, sort_array, struct
df.withColumn(
    "descending_columns", 
    array(
        *[
            sort_array(
                array(
                    *[
                        struct([col(c).alias("value"), lit(c).alias("key")]) 
                        for c in cols_to_sort
                    ]
                ), 
                asc=False
            )[i]["key"]
            for i in range(len(cols_to_sort))
        ]
    )
).show(truncate=False)
#+---+---+---+---+---+------------------+
#|ID |key|a  |b  |c  |descending_columns|
#+---+---+---+---+---+------------------+
#|0  |1  |5  |2  |1  |[a, b, c]         |
#|1  |1  |3  |4  |5  |[c, b, a]         |
#+---+---+---+---+---+------------------+

尽管这看起来很复杂，但它应该提供比udf 解决方案更好的性能。

更新：要在值相同的情况下按原始列顺序排序，您可以在包含索引的结构中插入另一个值。由于排序是降序的，我们使用索引的负数。

例如，如果您的输入数据框如下：

df.show()
#+---+---+---+---+---+
#| ID|key|  a|  b|  c|
#+---+---+---+---+---+
#|  0|  1|  5|  2|  1|
#|  1|  1|  3|  4|  5|
#|  2|  1|  4|  4|  5|
#+---+---+---+---+---+

上面最后一行的值在 a 和 b 之间。在这种情况下，我们希望a 排在b 之前。

df.withColumn(
    "descending_columns", 
    array(
        *[
            sort_array(
                array(
                    *[
                        struct(
                            [
                                col(c).alias("value"), 
                                lit(-j).alias("index"), 
                                lit(c).alias("key")
                            ]
                        ) 
                        for j, c in enumerate(cols_to_sort)
                    ]
                ), 
                asc=False
            )[i]["key"]
            for i in range(len(cols_to_sort))
        ]
    )
).show(truncate=False)
#+---+---+---+---+---+------------------+
#|ID |key|a  |b  |c  |descending_columns|
#+---+---+---+---+---+------------------+
#|0  |1  |5  |2  |1  |[a, b, c]         |
#|1  |1  |3  |4  |5  |[c, b, a]         |
#|2  |1  |4  |4  |5  |[c, a, b]         |
#+---+---+---+---+---+------------------+

【讨论】：

在平局的情况下，列名将作为第二个排序条件。
感谢 Pault 的回答。我的下一个挑战是在列条目相等时保持列名的顺序。你应该如何实现它？
您可以在结构中添加第二个条目，即列数组中的索引。然后，这将用作第一个决胜局。我可以发布更新。

【解决方案2】：

您可以将列插入单个结构并在 udf 中处理。

from pyspark.sql import functions as F
from pyspark.sql import types as T

name_cols = ['a', 'b', 'c']

def ordered_columns(row):
    return [x for _,x in sorted(zip(row.asDict().values(), name_cols), reverse=True)]
udf_ordered_columns = F.udf(ordered_columns, T.ArrayType(T.StringType()))

df1 = (
    df
    .withColumn(
        'row',
        F.struct(*name_cols)
    )
    .withColumn(
        'descending_columns',
        udf_ordered_columns('row')
    )
)

这样的东西应该可以工作，如果上面没有，那么请告诉我。

【讨论】：