PySpark处理ADLS Gen2中Base64文件：输出数量不符及分文件保存需求

问题原因 #

• 你看到的20个Pickle文件对应Spark RDD的分区数，而非输入文件数量。Spark读取文本文件时，会根据文件大小、集群默认并行度等参数拆分数据为多个分区，多个小文件可能被合并到同一个分区中，最终saveAsPickleFile会为每个分区生成一个文件，因此输出文件数和输入文件数不匹配。
• 当前代码没有保留输入文件的元信息，无法区分每条数据来自哪个源文件，自然无法实现“每个输入文件对应一个输出Pickle文件”的需求。

解决方案 #

要实现每个输入文件单独解码存储，需按以下步骤调整代码：

1. 读取数据时保留每行对应的源文件路径；
2. 按源文件路径分组，聚合同一文件的所有内容；
3. 遍历每个分组，单独解码并保存为Pickle文件。

修改后的代码示例 #

import base64
from datetime import datetime
from pyspark.sql.functions import input_file_name, collect_list

# 读取文本数据，同时添加源文件路径列
input_data = spark.read.text('/rawdata/no=*/*.txt').withColumn("source_file", input_file_name())

# 按源文件分组，收集该文件的所有行内容
grouped_data = input_data.groupBy("source_file").agg(collect_list("value").alias("file_content"))

# 定义单文件解码与保存函数
def process_single_file(row):
    source_file_path = row.source_file
    # 假设每个txt文件为单行Base64编码内容，若为多行则需先拼接内容再解码
    decoded_content = base64.b64decode(row.file_content[0])
    
    # 从源路径提取文件名，替换后缀为.pkl
    file_name = source_file_path.split("/")[-1].replace(".txt", ".pkl")
    timestamp = datetime.now().strftime("%Y%m%d%H%M")
    output_file_path = f"{output_file_name}_ss3_{timestamp}/{file_name}"
    
    # 写入解码后的内容到Pickle文件
    with open(output_file_path, "wb") as f:
        f.write(decoded_content)

# 遍历所有分组文件，逐个处理保存
# 若文件数量极大，建议改用foreachPartition避免Driver内存溢出
for row in grouped_data.collect():
    process_single_file(row)

注意事项 #

• 若txt文件包含多行Base64内容，需调整解码逻辑，例如先将file_content列表拼接为完整字符串再解码；
• 当输入文件数量极多时，collect()会将所有数据拉到Driver节点，可能引发内存溢出，此时建议使用grouped_data.foreachPartition(lambda partition: [process_single_file(row) for row in partition])，让Executor节点处理分区内的文件；
• 确保Spark作业对ADLS Gen2的输出路径拥有写入权限，路径需符合ADLS格式要求（如abfss://<container>@<account>.dfs.core.windows.net/...）。

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