PySpark 데이터 프레임에서 n 행에 대한 핵심 요소를 추가하십시오.

Question

아래와 같은 데이터 프레임이 있습니다.

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

나는 새 열이 다음 n 행의 처음 n 행에 대한 값 1과 2를해야합니다 Key 말을 추가하고 싶습니다. 결과는 다음과 같습니다 :

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)