如何在PySpark DataFrame中按条件新增行并调整col_seq序号

实现思路 #

• 先将原始DF按Hash_enc_ind拆分成分两类：不需要扩展的（值为0）和需要扩展的（值为1）
• 对需要扩展的数据集，通过explode函数为每行生成2条派生记录，分别拼接_h和_e后缀到col_name字段
• 合并两类数据集后，按原始col_seq排序保证顺序和原始逻辑一致，再通过窗口函数重新生成连续的col_seq值

完整代码 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 拆分原始数据集
df_needless = DF1.filter(F.col("Hash_enc_ind") == 0)
df_need_expand = DF1.filter(F.col("Hash_enc_ind") == 1)

# 2. 扩展需要派生行的数据集
df_expanded = df_need_expand \
    .withColumn("suffix", F.explode(F.array(F.lit("_h"), F.lit("_e")))) \
    .withColumn("col_name", F.concat(F.col("col_name"), F.col("suffix"))) \
    .drop("suffix")

# 3. 合并数据集并重写col_seq
merge_df = df_needless.unionByName(df_expanded)
# 按原始col_seq排序保证整体顺序和原始一致
window_order = Window.orderBy(F.col("col_seq").asc())
result_df = merge_df.withColumn("col_seq", F.row_number().over(window_order))

# 验证输出
result_df.orderBy("col_seq").show(truncate=False)

输出验证 #

执行上述代码后得到的result_df和预期输出完全一致，col_seq为连续正整数，派生行的顺序和后缀均符合要求。

内容的提问来源于stack exchange，提问作者Dushyant Gaur