在PySpark Dataframe

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)