处理排名关系：Pyspark

Question

我有一个像下面这样的数据框。

+---+-------+-------+
|ayy| artist|numbers|
+---+-------+-------+
|  a|  Monet|  10000|
|  a|   Dali|  10000|
|  a|Gauguin|  10000|
|  b|  Monet|  10000|
|  b|   Dali|  10000|
|  b|Gauguin|  10000|
+---+-------+-------+

我想根据“数字”选择一位“艺术家”，但我想确保当“数字”之间存在联系时我随机选择

我在下面实现了以下内容。

w = Window.partitionBy('ayy').orderBy(F.col('numbers').desc())
df_test = df_test.withColumn('rank', F.rank().over(w))\
                .withColumn('rank2', F.row_number().over(w))
df_test.show()

这给了我

+---+-------+-------+----+-----+
|ayy| artist|numbers|rank|rank2|
+---+-------+-------+----+-----+
|  a|  Monet|  10000|   1|    1|
|  a|Gauguin|  10000|   1|    2|
|  a|   Dali|  10000|   1|    3|
|  b|  Monet|  10000|   1|    1|
|  b|   Dali|  10000|   1|    2|
|  b|Gauguin|  10000|   1|    3|
+---+-------+-------+----+-----+

看，在这种情况下，我只能根据他们的“数字”找到一位艺术家，但我想确保在平局的情况下随机选择一位“艺术家”。

我的第一个想法是选择 udf，但我仍然不太确定如何去做。

Answer 1

您可以将 [0.0, 1.0) 范围内的随机数添加到排名中，然后根据排名分配行号：

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'rank', 
    F.rank().over(Window.partitionBy('ayy').orderBy(F.col('numbers').desc())) 
    + F.rand(seed=1)
).withColumn(
    'rank', 
    F.row_number().over(Window.partitionBy('ayy').orderBy('rank'))
)

df2.show()
+---+-------+-------+----+
|ayy| artist|numbers|rank|
+---+-------+-------+----+
|  b|Gauguin|  10000|   1|
|  b|   Dali|  10000|   2|
|  b|  Monet|  10000|   3|
|  a|  Monet|  10000|   1|
|  a|Gauguin|  10000|   2|
|  a|   Dali|  10000|   3|
+---+-------+-------+----+