PySpark/Databricks解析GS桶5M+XML文件的分片优化需求

我在GS桶中存储了超过500万个XML文件，想要用PySpark的多进程能力解析。但直接加载源目录时内存占用过高跑不起来，希望把加载解析过程按10个文件为一组分片执行，又不想逐个遍历文件，试了多种方法都没成功，求助。

当前使用的PySpark代码如下：

def read_process_files(path):
    spark = SparkSession.builder.getOrCreate()

    # Parsing the XML files
    df = spark.read.format('com.databricks.spark.xml').option('rowTag', 'article').load(path)

    # Apply the UDF to parse XML
    df_parsed = df.withColumn('parsed_data', parse_xml_udf(col('body')))

    # Add a boolean column verifying if the article is in the Pathway candidate list
    df_parsed = df_parsed.withColumn("criteria", col("parsed_data.issn").isin(candidate_issn_list))

    # Write to JSON files
    json_dir = f'gs://setup/processed/articles/'
    df_parsed.select(
        "parsed_data.article_title",
        "parsed_data.alt_title",
        "parsed_data.journal_title",
        "parsed_data.issn",
        "parsed_data.year",
        "parsed_data.abstract",
        "parsed_data.body",
        "parsed_data.uri",
        "parsed_data.doi",
        "criteria"
    ).write.mode("overwrite").format("json").option("ignoreNullFields", "false").save(json_dir)

read_process_files('gs://setup/raw/articles/xml_test*')

解决方案 #

方法1：基于RDD分组实现批量文件处理 #

利用Spark的RDD操作批量获取文件路径并分组，无需逐个遍历文件：

def read_process_files(path):
    spark = SparkSession.builder.getOrCreate()

    # 获取所有符合匹配规则的文件路径
    file_paths = spark.sparkContext.wholeTextFiles(path).keys().distinct()
    # 按每10个文件为一组进行分组
    grouped_paths = file_paths.zipWithIndex()\
        .groupBy(lambda x: x[1] // 10)\
        .map(lambda x: [p[0] for p in x[1]])

    # 遍历每个分组处理
    for group in grouped_paths.collect():
        # 加载当前分组的文件
        df = spark.read.format('com.databricks.spark.xml')\
            .option('rowTag', 'article')\
            .load(','.join(group))
        
        # 原解析逻辑保留
        df_parsed = df.withColumn('parsed_data', parse_xml_udf(col('body')))
        # 广播ISSN列表减少内存开销
        broadcast_issn = spark.sparkContext.broadcast(candidate_issn_list)
        df_parsed = df_parsed.withColumn("criteria", col("parsed_data.issn").isin(broadcast_issn.value))
        
        # 改用append模式避免覆盖已处理结果
        json_dir = f'gs://setup/processed/articles/'
        df_parsed.select(
            "parsed_data.article_title",
            "parsed_data.alt_title",
            "parsed_data.journal_title",
            "parsed_data.issn",
            "parsed_data.year",
            "parsed_data.abstract",
            "parsed_data.body",
            "parsed_data.uri",
            "parsed_data.doi",
            "criteria"
        ).write.mode("append").format("json").option("ignoreNullFields", "false").save(json_dir)

read_process_files('gs://setup/raw/articles/xml_test*')

方法2：调整Spark参数控制单批次处理文件数 #

通过配置Spark的文件读取参数，让框架自动控制每个分区的文件数量，无需手动分组：

def read_process_files(path):
    # 配置参数：根据单个XML文件大小调整maxPartitionBytes，确保每个分区约包含10个文件
    spark = SparkSession.builder\
        .config("spark.sql.files.maxPartitionBytes", "64MB")\  # 示例值，需根据实际文件大小调整
        .config("spark.sql.files.openCostInBytes", "4MB")\    # 提升小文件的合并优先级
        .getOrCreate()

    # 原加载逻辑不变，Spark会自动按配置合并文件到分区
    df = spark.read.format('com.databricks.spark.xml').option('rowTag', 'article').load(path)

    df_parsed = df.withColumn('parsed_data', parse_xml_udf(col('body')))
    broadcast_issn = spark.sparkContext.broadcast(candidate_issn_list)
    df_parsed = df_parsed.withColumn("criteria", col("parsed_data.issn").isin(broadcast_issn.value))

    json_dir = f'gs://setup/processed/articles/'
    df_parsed.select(
        "parsed_data.article_title",
        "parsed_data.alt_title",
        "parsed_data.journal_title",
        "parsed_data.issn",
        "parsed_data.year",
        "parsed_data.abstract",
        "parsed_data.body",
        "parsed_data.uri",
        "parsed_data.doi",
        "criteria"
    ).write.mode("overwrite").format("json").option("ignoreNullFields", "false").save(json_dir)

read_process_files('gs://setup/raw/articles/xml_test*')

方法3：用结构化流实现批量触发处理 #

如果允许用流式处理模式，可以通过maxFilesPerTrigger参数限制每次处理的文件数：

def read_process_files(path):
    spark = SparkSession.builder.getOrCreate()

    # 流式读取，每次触发处理10个文件
    df = spark.readStream.format('com.databricks.spark.xml')\
        .option('rowTag', 'article')\
        .option('maxFilesPerTrigger', 10)\
        .load(path)

    df_parsed = df.withColumn('parsed_data', parse_xml_udf(col('body')))
    broadcast_issn = spark.sparkContext.broadcast(candidate_issn_list)
    df_parsed = df_parsed.withColumn("criteria", col("parsed_data.issn").isin(broadcast_issn.value))

    # 启动流处理，写入结果
    query = df_parsed.select(
        "parsed_data.article_title",
        "parsed_data.alt_title",
        "parsed_data.journal_title",
        "parsed_data.issn",
        "parsed_data.year",
        "parsed_data.abstract",
        "parsed_data.body",
        "parsed_data.uri",
        "parsed_data.doi",
        "criteria"
    ).writeStream\
        .format("json")\
        .option("ignoreNullFields", "false")\
        .option("path", "gs://setup/processed/articles/")\
        .option("checkpointLocation", "gs://setup/processed/checkpoint/")\  # 必须设置检查点目录
        .start()

    query.awaitTermination()  # 等待所有文件处理完成

read_process_files('gs://setup/raw/articles/xml_test*')

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