时间戳

首页 > 解决方案 > 分批拆分数据帧pyspark

pyspark - 分批拆分数据帧pyspark

问题描述

我的要求是将数据帧分成 2 个批次，每个批次仅包含 2 个项目，并且批次大小（输出中的 BATCH）逐渐增加。

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

输出/输出

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

标签： pyspark

解决方案

我能够通过以下方法实现这一目标：

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)

推荐阅读