Pyspark使用wholeTextFiles遭遇java.lang.OutOfMemoryError问题求助

问题描述 #

有1160个单文件大小300MB的XML文件（总数据量约300GB），在256核、256GB内存的本地机器上用Spark统计总词数，调整内存配置和CPU核心数后仍报java.lang.OutOfMemoryError: Java heap space错误，两次测试代码及报错如下：

初始测试代码 #

from pyspark import SparkConf, SparkContext

if __name__ == "__main__":
    data_dir = '/shared/hm31/xml_data/'

    conf = SparkConf().setMaster("local[*]").setAppName("WordCount")
    conf.set("spark.driver.memory", "200g")
    conf.set("spark.executor.memory", "6g")  # 尝试增大executor内存

    sc = SparkContext(conf = conf)

    rdd = sc.wholeTextFiles(data_dir,500)
    words = rdd.flatMap(lambda x: x.split())
    print("words",words.count())

更新后的测试代码 #

data_dir = '/shared/hm31/xml_data/'

conf = SparkConf().setMaster("local[3]").setAppName("WordCount")
conf.set("spark.driver.memory", "200g")

sc = SparkContext(conf = conf)

rdd = sc.wholeTextFiles(data_dir,512)
words = rdd.flatMap(lambda x: x.split())
print("Words",words.count())

核心报错信息 #

java.lang.OutOfMemoryError: Java heap space
        at org.sparkproject.guava.io.ByteStreams.toByteArray(ByteStreams.java:252)
        at org.apache.spark.input.WholeTextFileRecordReader.nextKeyValue(WholeTextFileRecordReader.scala:79)

问题根源 #

1. wholeTextFiles的内存开销：该API会在Java层把整个文件内容读成字节数组，再转成字符串传给Python进程，单个300MB的文件转成字符串后内存占用会翻倍以上，若多个文件合并到一个分区，内存直接溢出。
2. 分区数设置不合理：之前设置的分区数（500/512）少于文件总数（1160），导致每个分区要处理多个文件，进一步加剧内存压力。
3. 本地模式内存配置遗漏：本地模式下Executor和Driver是同一个进程，spark.executor.memory配置无效，且未设置spark.driver.memoryOverhead（用于非堆内存），导致堆内存被占满。

解决方案 #

方案一：自定义文件读取逻辑，控制分区粒度 #

直接并行化文件路径，让每个分区仅处理单个文件，避免Spark自动合并文件，同时绕开wholeTextFiles的内存瓶颈：

from pyspark import SparkConf, SparkContext
import os

if __name__ == "__main__":
    data_dir = '/shared/hm31/xml_data/'

    # 本地模式下用64核（避免核心过多导致内存竞争），调整内存参数
    conf = SparkConf().setMaster("local[64]").setAppName("XMLWordCount")
    conf.set("spark.driver.memory", "100g")  # 分配100G堆内存给Driver
    conf.set("spark.driver.memoryOverhead", "50g")  # 分配50G非堆内存
    # 强制每个Split对应一个300MB的文件，避免Hadoop拆分/合并
    conf.set("spark.hadoop.mapreduce.input.fileinputformat.split.maxsize", "314572800")
    conf.set("spark.hadoop.mapreduce.input.fileinputformat.split.minsize", "314572800")

    sc = SparkContext(conf=conf)

    # 获取所有XML文件的完整路径
    file_list = [os.path.join(data_dir, filename) for filename in os.listdir(data_dir) 
                 if filename.lower().endswith('.xml')]
    
    # 并行化文件路径，分区数等于文件数，确保每个分区处理一个文件
    files_rdd = sc.parallelize(file_list, numSlices=len(file_list))

    # 定义单个文件的词数统计函数
    def count_words(file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            return len(content.split())

    # 计算每个文件的词数，再累加总和
    total_words = files_rdd.map(count_words).sum()

    print(f"总词数：{total_words}")
    sc.stop()

方案二：使用Spark XML专用库（更高效） #

如果允许添加依赖，使用Databricks的Spark XML库可以按XML节点读取数据，避免加载整个文件：

1. 添加依赖：启动Spark时加上--packages com.databricks:spark-xml_2.12:0.15.0（根据Spark版本调整）
2. 示例代码：

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .master("local[64]") \
    .appName("XMLWordCount") \
    .config("spark.driver.memory", "100g") \
    .config("spark.driver.memoryOverhead", "50g") \
    .getOrCreate()

# 读取XML文件，按文本内容字段统计词数
df = spark.read.format("xml") \
    .option("rowTag", "your_root_tag")  # 替换为你的XML根节点标签
    .load("/shared/hm31/xml_data/")

# 统计总词数
total_words = df.rdd.flatMap(lambda row: row[0].split()).count()
print(f"总词数：{total_words}")

spark.stop()

关键优化点说明 #

• 分区粒度控制：让每个分区对应单个文件，避免多文件合并导致的内存过载。
• 绕开wholeTextFiles：直接用Python原生IO读取文件，减少Java-Python之间的数据传输开销。
• 内存参数调整：本地模式下重点配置spark.driver.memory和spark.driver.memoryOverhead，而非spark.executor.memory。

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