处理排名中的关系: 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|
+---+-------+-------+----+