为 PySpark 数据帧中的 n 行添加一个键元素

Question

>我有一个如下所示的数据帧。

id   |  run_id
--------------
4    |  12345
6    |  12567
10   |  12890
13   |  12450

我想添加一个新列，比如说Key前 n 行的值为 1，接下来的 n 行的值为 2。结果将是这样的：

id   |  run_id  | key
----------------------
4    |  12345   | 1
6    |  12567   | 1
10   |  12890   | 2
13   |  12450   | 2

是否有可能对 PySpark 做同样的事情？提前感谢您的帮助。

Answer 1

这是使用 zipWithIndex 执行此操作的一种方法：

# sample rdd
rdd=sc.parallelize([[4,12345], [6,12567], [10,12890], [13,12450]])
# group size for key    
n=2
# add rownumber and then label in batches of size n
rdd=rdd.zipWithIndex().map(lambda (x, rownum): x+[int(rownum/n)+1])
# convert to dataframe
df=rdd.toDF(schema=['id', 'run_id', 'key'])
df.show(4)