Question

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

+---+---+---+---+---+
| 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版。

Answer 1

您可以将列插入单个结构并在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')
    )
)

类似的东西应该可以工作，如果上面的方法不可行，请告诉我。

Answer 2

对于Spark版本<2.4，您可以使用Confluent Docs和sort_array，而无需udf来实现。

首先获取要排序的列的列表

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]         |
#+---+---+---+---+---+------------------+

Spark DataFrame：基于列值获取按行排序的列名称

2 个答案: