Question

我有以下PySpark DataFrame：

+----+----------+----------+----------+
|  id|         a|         b|         c|
+----+----------+----------+----------+
|2346|2017-05-26|      null|2016-12-18|
|5678|2013-05-07|2018-05-12|      null|
+----+----------+----------+----------+

我的理想输出是：

+----+---+---+---+
|id  |a  |b  |c  |
+----+---+---+---+
|2346|2  |0  |1  |
|5678|1  |2  |0  |
+----+---+---+---+

该行中的日期越近，得分越高

我看过类似的文章，建议使用窗口功能。问题是我需要在行而不是列中对值进行排序。

Answer 1

您可以将每一行中的值放入数组中，并使用pyspark.sql.functions.sort_array()对其进行排序。

import pyspark.sql.functions as f
cols = ["a", "b", "c"]
df = df.select("*", f.sort_array(f.array([f.col(c) for c in cols])).alias("sorted"))
df.show(truncate=False)
#+----+----------+----------+----------+------------------------------+
#|id  |a         |b         |c         |sorted                        |
#+----+----------+----------+----------+------------------------------+
#|2346|2017-05-26|null      |2016-12-18|[null, 2016-12-18, 2017-05-26]|
#|5678|2013-05-07|2018-05-12|null      |[null, 2013-05-07, 2018-05-12]|
#+----+----------+----------+----------+------------------------------+

现在，您可以结合使用pyspark.sql.functions.coalesce()和pyspark.sql.functions.when()来遍历cols中的每一列，并在排序后的数组中找到相应的索引。

df = df.select(
    "id",
    *[
        f.coalesce(
            *[
                f.when(
                    f.col("sorted").getItem(i) == f.col(c),
                    f.lit(i)
                ) 
                for i in range(len(cols))
            ]
        ).alias(c) 
        for c in cols
    ]
)
df.show(truncate=False)
#+----+---+----+----+
#|id  |a  |b   |c   |
#+----+---+----+----+
#|2346|2  |null|1   |
#|5678|1  |2   |null|
#+----+---+----+----+

最后用0填充null值：

df = df.na.fill(0)
df.show(truncate=False)
#+----+---+---+---+
#|id  |a  |b  |c  |
#+----+---+---+---+
#|2346|2  |0  |1  |
#|5678|1  |2  |0  |
#+----+---+---+---+

在行内对DataFrame进行排序并获得排名

1 个答案: