在Pyspark中为1到N之间的行分配一个随机数

Question

我正在尝试找出最好的方法，将1到N之间的数字随机分配给一行，以使每一行都是不同的。 （N是数据集中的行数）。行数可能超过10M，因此可能不是这样。

最初的想法是使用udf生成数组，但是我不确定如何将数组的每个索引分配给一行。

def create_rand_range(end):
    return list(random.sample(range(1, end), end-1))

示例： n = 3 create_rand_range = [3,1,2]

数据表输出：

| Col 1 | Rand_Output |
|-------|-------------|
| A     | 3           |
| B     | 1           |
| C     | 2           |

任何帮助都会很棒。

Answer 1

下面是几个为行分配编号的示例。

选项1：使用monotonically_increasing_id()函数。

##sample dataframe    
df = spark.createDataFrame([('A',),('B',),('C',),('D',) ],["dummy"])

##monotonically_increasing_id function
from pyspark.sql import functions as psf
df.withColumn("Rand_Output ",psf.monotonically_increasing_id()).show()


#+-----+------------+
#|dummy|Rand_Output |
#+-----+------------+
#|    A|  8589934592|
#|    B| 25769803776|
#|    C| 42949672960|
#|    D| 60129542144|
#+-----+------------+

选项2：使用Window函数

from pyspark.sql import functions as psf
from pyspark.sql import window as psw

w = psw.Window().partitionBy(psf.lit('a')).orderBy(psf.lit('a'))
df.withColumn("row_num", psf.row_number().over(w)).show()

#+-----+-------+
#|dummy|row_num|
#+-----+-------+
#|    A|      1|
#|    B|      2|
#|    C|      3|
#|    D|      4|
#+-----+-------+

Answer 2

您可以通过udf这样做：

import numpy as np
import random
from pyspark.sql.types import IntegerType
import pyspark.sql.functions as F

num_list = np.arange(1, df.count()+1, 1).tolist()

def delete_rand_items():
    global num_list
    to_delete = random.sample(range(len(num_list)),1)
    x = num_list[to_delete[0]]
    num_list = [x for i,x in enumerate(num_list) if not i in to_delete]
    return x

，然后使用此udf定义新列：

assign_num = F.udf(delete_rand_items, IntegerType())
df = df.withColumn('Rand_output', assign_num())