如何用Spark Python并行读取目录文件并提取文件名与首行

这问题我刚好处理过，用Spark来并行提取每个文件的文件名和首行有两种常用方案，你可以根据文件的大小和业务场景选择：

方法一：使用RDD的wholeTextFiles API（适合小文件批量处理） #

这种方法会把每个文件的完整内容读取为一个字符串，直接提取首行，逻辑简单且能保证首行的准确性，适合文件数量多但单个文件体积不大的场景。

from pyspark import SparkContext

# 初始化SparkContext（如果已经有上下文可以跳过这步）
sc = SparkContext("local[*]", "ExtractFileHeader")

# 读取目标目录下的所有文件，返回格式为 RDD[(文件路径, 文件完整内容)]
file_content_rdd = sc.wholeTextFiles("/path/to/your/source/directory")

# 提取文件名和首行：分割内容为行，取第一行；处理无换行的文件
header_rdd = file_content_rdd.map(
    lambda x: (
        x[0], 
        x[1].split('\n')[0] if '\n' in x[1] else x[1]
    )
)

# 收集结果（如果文件数量极大，建议写入存储而非collect）
results = header_rdd.collect()

# 打印或处理结果
for file_path, first_line in results:
    print(f"文件: {file_path} | 首行: {first_line}")

注意点： #

• 如果是本地文件系统，路径需要加上file://前缀（例如file:///home/user/data）；HDFS路径直接写hdfs://namenode:port/path即可。
• 若单个文件体积过大，这种方法会占用较多内存，不建议使用。

方法二：使用DataFrame API（适合大文件场景） #

如果你的文件体积较大，用行级读取的DataFrame方式更节省内存，结合Spark内置的input_file_name()函数获取文件名，再通过窗口函数筛选每个文件的第一行。

from pyspark.sql import SparkSession
from pyspark.sql.functions import input_file_name, row_number
from pyspark.sql.window import Window

# 初始化SparkSession
spark = SparkSession.builder.appName("ExtractFileHeader").getOrCreate()

# 读取目录下所有文件，每行作为一条记录
text_df = spark.read.text("/path/to/your/source/directory")

# 添加文件名列
df_with_filename = text_df.withColumn("file_path", input_file_name())

# 按文件名分区，给每个文件内的行编号
window_spec = Window.partitionBy("file_path").orderBy("value")
df_with_row_num = df_with_filename.withColumn("row_num", row_number().over(window_spec))

# 筛选出每个文件的第一行（行号=1）
header_df = df_with_row_num.filter(df_with_row_num.row_num == 1).select("file_path", "value")

# 查看结果或写入存储
header_df.show(truncate=False)
# header_df.write.mode("overwrite").csv("/path/to/output")

注意点： #

• 窗口函数的orderBy这里只是为了满足语法要求，如果需要严格保证是文件的物理首行，Spark读取文本文件时会保留行的原始顺序；如果文件被拆分到多个分区，一般场景下默认顺序足够覆盖需求。
• 这种方法支持大文件的分布式处理，内存压力小。

内容的提问来源于stack exchange，提问作者A.N.Gupta