基于权重列重复Spark DataFrame行的过采样优化方案求助

针对你遇到的Pandas过采样效率低、内存溢出的问题，我给你提供一个基于Spark DataFrame的高效分布式解决方案，完全能利用你的15个CPU核心，再也不会出现内存耗尽的情况！

核心思路 #

Spark的优势在于分布式计算，我们可以利用Spark内置的高阶函数sequence生成对应重复次数的序列，再通过explode展开序列来实现行的重复，整个过程完全分布式执行，充分利用集群资源。

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode, sequence, lit, col, round

# 初始化Spark会话
spark = SparkSession.builder.appName("OversampleByWeight").getOrCreate()

# 读取CSV数据，注意处理分隔符和首尾空格
df = spark.read.csv("file.csv", sep="|", header=True, trim=True)

# 预处理Weight列：如果是浮点数先四舍五入，再转整数；同时过滤无效权重（<=0）
df = df.withColumn("Weight", round(col("Weight")).cast("int")) \
       .filter(col("Weight") > 0)

# 生成重复序列并展开，实现过采样
oversampled_df = df.withColumn("repeat_sequence", sequence(lit(1), col("Weight"))) \
                   .withColumn("dummy", explode(col("repeat_sequence"))) \
                   .drop("Weight", "repeat_sequence", "dummy")

# 查看结果（数据量大时可加limit限制行数）
oversampled_df.show()

# 可选：将结果保存回CSV
oversampled_df.write.csv("oversampled_result.csv", sep="|", header=True, mode="overwrite")

为什么这个方案高效？ #

• 分布式执行：Spark会将数据分片到多个Executor节点，同时调用你所有的15个CPU核心，彻底告别Pandas单线程的低效瓶颈
• 内存友好：不需要将全量数据加载到单节点内存，每个节点只处理部分数据，完全避免65G内存耗尽的崩溃问题
• 内置优化：sequence和explode都是Spark官方优化过的高阶函数，执行效率远高于Pandas+NumPy的单机方案

注意事项 #

1. 如果你的Weight列已经是整数，可以去掉round函数，直接转int即可
2. 过滤Weight>0是为了避免生成空序列导致的错误，如果你的数据中没有权重为0或负数的行，可以省略这一步
3. 保存结果时，mode="overwrite"会覆盖已有文件，根据你的需求可以调整为append或其他模式

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