如何在使用pyspark从s3读取文件时读取文件名

Question

我有一个用例，在 s3 中有多个文件，我正在阅读如下：

df1 = spark.read.csv("s3://bucket/fact/*.dat")

所有 .dat 文件的开头都有 6 位数字，即 PO id

190234_purcahse.dat
125134_purcahse.dat

读取时我需要在数据帧 df1 中作为新列的这个 PO id。我怎样才能以最有效的方式实现这一目标？ 有什么方法可以在读取文件时获取文件名？

Answer 1

使用来自 spark 的 input_file_name() 函数。

df1 = spark.read.csv("s3://bucket/fact/*.dat").withColumn("fn",input_file_name())

然后使用 regexp_extract() 函数提取您的 PO id。