Блог

Мое требование состоит в том, чтобы разделить фрейм данных на группу из 2 пакетов, при этом каждый пакет содержит только 2 элемента, а размер пакета (ПАКЕТ на выходе) увеличивается постепенно.

 col#1 col#2 DATE
A     1     202010
B     1.1   202010
C     1.2   202010 
D     1.3   202001
E     1.4   202001
  

O/P

 col#1 col#2 DATE    BATCH
A     1     202010  1
B     1.1   202010  1
C     1.2   202010  2
D     1.3   202001  2
E     1.4   202001  3 
  

Я смог добиться этого, следуя следующему подходу:

 def dfZipWithIndex (df, offset=1, colName='rowId'):   
   new_schema = StructType([StructField(colName,LongType(),True)]  
   df.schema.fields)   
   zipped_rdd = df.rdd.zipWithIndex()   
   new_rdd =zipped_rdd.map(lambda args: ([args[1]   offset]   list(args[0])))   
   return spark.createDataFrame(new_rdd, new_schema)

chunk_size=2 
final_new=dfZipWithIndex(input_df) 


temp_final=input_df.withColumn('BATCH',F.floor(F.col('rowId')/chunk_size) 1)