使用重复次数相等的范围内的数字填充 Pyspark 数据帧列

Question

我想用一系列数字填充pyspark数据帧上的新列，根据df的大小，该列重复大约相同的次数。 现在，我正在使用以下代码生成 1 到 3 之间的随机数并填充新列。问题是每个数字重复的次数不相等。

df = df1.withColumn("random", func.round(rand()*(3-1)+1,0).cast('integer'))

Answer 1

与其在整个数据集中应用random，这确实是随机的，并且您对其分布的控制较少，不如在随机排序数据集时使用窗口函数row_number

from pyspark.sql import functions as F
from pyspark.sql import Window as W
(df
.withColumn('_id', F.row_number().over(W.orderBy(F.rand())))
.withColumn('_grp', F.col('_id') % 3 + 1)
.show()
)
+---+---+---+---+---+----+
|  a|  b|  c|  d|_id|_grp|
+---+---+---+---+---+----+
|  A|  B|  C|  D|  1|   2|
|  A|  B|  C|  D|  2|   3|
|  B|  C|  D|  A|  3|   1|
|  B|  C|  D|  E|  4|   2|
|  B|  C|  D|  A|  5|   3|
|  B|  C|  D|  E|  6|   1|
|  B|  C|  D|  E|  7|   2|
|  B|  C|  D|  E|  8|   3|
|  B|  C|  D|  A|  9|   1|
|  B|  C|  D|  A| 10|   2|
|  A|  B|  C|  D| 11|   3|
|  A|  B|  C|  D| 12|   1|
|  B|  C|  D|  E| 13|   2|
+---+---+---+---+---+----+