如何将目录下所有TXT文件读取为带文件名的Pandas DataFrame

当前现状 #

你现在的代码只能读取单个指定TXT文件，文件内容是空格分隔的无表头数据，示例如下：

0 0.712518 0.615250 0.439180 0.206500
1 0.635078 0.811750 0.292786 0.092500

原始单文件处理代码：

spark.read.format('csv').options(header='false').load("/mnt/datasets/model1/train/labels/2a.txt").toPandas()
df_2a.columns = ['Value']
df_2a_split = df_2a['Value'].str.split(' ', n=0, expand=True)
df_2a_split.columns = ['class','c1','c2','c3','c4']
display(df_2a_split)

现在需要实现：读取目录下所有TXT文件，把文件名作为DataFrame的首列，最终期望输出格式如下：

file_name class   c1       c2      c3          c4
2a.txt  0   0.712518    0.61525 0.43918     0.2065
2a.txt  1   0.635078    0.81175 0.292786    0.0925
2b.txt  2   0.551273    0.5705  0.30198     0.0922
2b.txt  0   0.550212    0.31125 0.486563    0.2455

改进后的代码 #

直接用Spark读取整个目录，同时获取文件名，一步到位，不用再用Pandas拆分列，效率更高：

from pyspark.sql.functions import input_file_name, substring_index

# 读取目录下所有TXT文件，指定空格为分隔符，自动推断列类型
df_spark = spark.read.format('csv') \
    .options(header='false', sep=' ', inferSchema=True) \
    .load("/mnt/datasets/model1/train/labels/")

# 添加文件名列：从完整文件路径中截取最后一段作为文件名
df_spark_with_file = df_spark.withColumn(
    'file_name',
    substring_index(input_file_name(), '/', -1)
)

# 设置列名并调整顺序，把file_name放在第一列
column_names = ['file_name', 'class', 'c1', 'c2', 'c3', 'c4']
df_spark_with_file = df_spark_with_file.toDF(*column_names)

# 转成Pandas DataFrame并展示
df_final = df_spark_with_file.toPandas()
display(df_final)

代码说明 #

• 直接读取目录路径，Spark会自动加载目录下所有TXT文件
• 用input_file_name()获取每行对应的完整文件路径，再通过substring_index截取文件名部分
• 提前指定分隔符为空格，Spark会直接把每行拆分成多列，省去后续Pandas拆分的步骤
• 通过toDF()一次性设置列名并调整顺序，确保file_name在首列

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