如何在PySpark中实现GroupBy Count的逆操作？

需求说明 #

输入DataFrame包含班级、男性计数、女性计数三列，需要将计数展开为对应行数的记录：男性计数对应性别'm'，女性计数对应性别'f'，最终得到每个班级下每条记录对应单个性别的结果。

输入DataFrame构造 #

先模拟创建输入数据：

from pyspark.sql import SparkSession
from pyspark.sql.functions import array_repeat, explode, col

spark = SparkSession.builder.appName("count_unfold").getOrCreate()

data = [
    ("A", 2, 1),
    ("B", 3, 1),
    ("C", 0, 3),
    ("D", 2, 4)
]
input_df = spark.createDataFrame(data, ["班级（class）", "男性计数（malecount）", "女性计数（femalecount）"])

方法一：拆分男女数据分别展开后合并 #

这种方法逻辑直观，分别处理男性和女性的计数展开，再合并结果：

# 处理男性记录：过滤计数>0的行，生成对应数量的'm'数组并展开
male_df = input_df.filter(col("男性计数（malecount）") > 0) \
                  .select(
                      col("班级（class）").alias("班级（Class）"),
                      explode(array_repeat("m", col("男性计数（malecount）"))).alias("性别（Gender）")
                  )

# 处理女性记录：同理处理女性计数
female_df = input_df.filter(col("女性计数（femalecount）") > 0) \
                    .select(
                        col("班级（class）").alias("班级（Class）"),
                        explode(array_repeat("f", col("女性计数（femalecount）"))).alias("性别（Gender）")
                    )

# 合并两个结果DF
output_df = male_df.union(female_df)

# 查看结果（可选）
output_df.show()

方法二：先宽表转长表再统一展开 #

如果后续有更多性别列需要处理，这种方法扩展性更好，先把宽表转为长表（性别-计数的键值对），再统一展开：

from pyspark.sql.functions import expr

# 宽表转长表：将男女计数列转为性别和计数列
melted_df = input_df.select(
    col("班级（class）"),
    expr("stack(2, 'm', 男性计数（malecount）, 'f', 女性计数（femalecount）) as (性别（Gender）, 计数)")
).filter(col("计数") > 0)

# 展开计数为对应行数
output_df = melted_df.select(
    col("班级（class）").alias("班级（Class）"),
    explode(array_repeat(col("性别（Gender）"), col("计数"))).alias("性别（Gender）")
)

# 查看结果（可选）
output_df.show()

两种方法都能得到预期的输出结果，方法二更适合多性别列的场景，方法一逻辑更直白易理解。

内容的提问来源于stack exchange，提问作者pradeep nadarajan