PySpark保存500万行DataFrame时Java堆内存溢出问题求助

一、换用低内存占用的存储格式 #

既然Parquet写入触发内存溢出，试试这些兼容性强、内存压力更小的格式：

• CSV格式：写入时限制单文件记录数，避免大文件拖垮内存，其他脚本用pandas/PySpark都能直接读：

  stream_df.write.mode("overwrite").option("maxRecordsPerFile", 100000).csv("./Stream3_csv")
  

• JSON格式：同样支持分块写入，结构化数据读取无压力：

  stream_df.write.mode("overwrite").option("maxRecordsPerFile", 80000).json("./Stream3_json")
  

• ORC格式：列式存储但内存占用比Parquet略低，稳定性更好：

  stream_df.write.mode("overwrite").orc("./Stream3_orc")
  

二、优化数据处理流程，从根源减内存开销 #

你拆分时间后再关联主表的操作会产生大量shuffle，试试直接拆分合并，跳过关联步骤：

# 直接拆分出开始事件记录
start_df = df.select(
    col("ID"),
    (col("Start_Time").cast("timestamp").cast("long") / RUNTIME).alias("Stream_Time"),
    lit("START").alias("Event_Type"),
    # 只选后续脚本需要的列，别带全量字段
    col("Severity"), col("Weather_Condition"), col("State")
)

# 拆分出结束事件记录
end_df = df.select(
    col("ID"),
    (col("End_Time").cast("timestamp").cast("long") / RUNTIME).alias("Stream_Time"),
    lit("END").alias("Event_Type"),
    col("Severity"), col("Weather_Condition"), col("State")
)

# 合并两个DataFrame，直接得到最终的stream_df，不用再关联
stream_df = start_df.union(end_df)

这样省掉了ID关联的shuffle操作，内存占用会大幅降低。另外记得只保留必要字段，别把没用的列都带上。

三、调整Spark配置的关键参数 #

除了spark.driver.memory，这些参数能帮你避开OOM：

from pyspark.sql import SparkSession
spark = SparkSession.builder \
    .appName("AccidentStream") \
    .config("spark.driver.memory", "8g")  # 别设太大，避免终端/Jupyter崩溃
    .config("spark.executor.memory", "6g") \
    .config("spark.driver.maxResultSize", "4g")  # 限制Driver端结果的最大尺寸
    .config("spark.sql.shuffle.partitions", "150")  # 根据数据量调shuffle分区数，默认200可适当减
    .getOrCreate()

# 开启内存溢出时自动写磁盘
spark.conf.set("spark.storage.memoryFraction", "0.5")
spark.conf.set("spark.sql.inMemoryColumnarStorage.compressed", "true")

四、分批次写入数据 #

如果数据量实在太大，一次性写入风险高，就分批次处理：

# 获取所有唯一ID，分批处理
id_list = df.select("ID").distinct().rdd.map(lambda x: x[0]).collect()
batch_size = 10000  # 每批处理10000个ID

for idx in range(0, len(id_list), batch_size):
    batch_ids = id_list[idx:idx+batch_size]
    # 过滤当前批次的数据
    batch_data = df.filter(col("ID").isin(batch_ids))
    # 执行拆分逻辑
    start_batch = batch_data.select(...)
    end_batch = batch_data.select(...)
    stream_batch = start_batch.union(end_batch)
    # 追加写入到目标路径
    stream_batch.write.mode("append").csv("./Stream3_csv_batched")

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