如何在PySpark中获取当前作业生成的输出文件名？

当使用PySpark将DataFrame写入指定目录（代码示例如下）时，默认生成part-****格式的文件，若存在多作业写入同一目录的场景，可通过以下几种方式获取本次作业生成的具体文件名：

df.write \
    .format(file_type) \
    .option(strings.str_header_lc, header) \
    .option(strings.str_delimiter_lc, delimiter) \
    .mode(mode) \
    .save(path)

方法一：对比写入前后的目录文件列表 #

1. 写入操作前，先获取目标目录下已存在的所有文件列表
2. 执行DataFrame写入操作
3. 再次获取目标目录的文件列表，两次列表的差异部分即为本次作业生成的文件

示例代码：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("GetNewPartFiles").getOrCreate()

# 目标路径
target_path = "/your/target/directory"

# 获取目录下所有文件的工具函数
def list_directory_files(spark, path):
    file_paths = spark.sparkContext.wholeTextFiles(f"{path}/*").keys().collect()
    return set(file_paths)

# 记录写入前的文件
pre_write_files = list_directory_files(spark, target_path)

# 执行写入操作
df.write \
    .format(file_type) \
    .option(strings.str_header_lc, header) \
    .option(strings.str_delimiter_lc, delimiter) \
    .mode(mode) \
    .save(target_path)

# 记录写入后的文件
post_write_files = list_directory_files(spark, target_path)

# 计算差异得到本次生成的文件
new_generated_files = post_write_files - pre_write_files
print("本次作业生成的文件:", new_generated_files)

方法二：自定义输出文件名（推荐小数据集使用） #

如果业务允许，直接指定输出文件名可避免后续对比操作。注意：当DataFrame分区数大于1时，需先合并为单个分区（coalesce(1)），但这会将数据集中到单个节点，仅适合小数据量场景。

示例代码：

import glob
import shutil

# 先写入临时目录
temp_dir = f"{target_path}/temp_output"
df.coalesce(1) \
    .write \
    .format(file_type) \
    .option(strings.str_header_lc, header) \
    .option(strings.str_delimiter_lc, delimiter) \
    .mode(mode) \
    .save(temp_dir)

# 找到生成的part文件并重命名为自定义名称
part_file = glob.glob(f"{temp_dir}/part-*")[0]
custom_file_name = f"{target_path}/my_custom_file.csv"
shutil.move(part_file, custom_file_name)

# 删除临时目录
shutil.rmtree(temp_dir)

方法三：利用Spark监听事件捕获输出文件 #

通过自定义SparkListener监听作业结束事件，从事件中提取输出文件信息。这种方式适合复杂场景，但依赖Spark内部API，不同版本可能存在兼容性问题。

示例代码片段：

from pyspark import SparkContext
from pyspark.scheduler import SparkListener

class OutputFileListener(SparkListener):
    def onJobEnd(self, job_end_event):
        job_id = job_end_event.jobId
        # 结合Spark内部API获取该作业对应的输出文件路径
        # 不同Spark版本实现逻辑可能不同，需自行适配
        pass

# 注册监听
sc = SparkContext.getOrCreate()
sc.addSparkListener(OutputFileListener())

# 执行写入操作
df.write.save(target_path)

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