从多个S3存储桶导入pyspark数据帧，其中一列表示该条目来自哪个存储桶

Question

我有一个按日期划分的S3存储桶列表。第一个存储桶的名称为2019-12-1，第二个存储桶的名称为2019-12-2，等等。

每个存储桶都将我正在读取的实木复合地板文件存储到pyspark数据帧中。从每个存储桶生成的pyspark数据帧具有完全相同的架构。我想做的是遍历这些存储桶，并将所有这些镶木地板文件存储到单个pyspark数据帧中，该数据帧具有一个日期列，该日期列指示该数据帧中每个条目实际来自哪个存储桶。

由于分别导入每个存储桶时生成的数据帧的架构深很多层（即，每一行包含结构数组的结构等），我想将所有存储桶组合到一个数据帧中的唯一方法是拥有一个具有单个“日期”列的数据框。 “日期”列的每一行将保存该日期相应S3存储桶的内容。

我可以在此行中读取所有日期：

df = spark.read.parquet("s3://my_bucket/*")

我已经看到有人通过在此行上添加“ withColumn”调用以创建“日期”列来达到我的描述，但是我不记得怎么做。

Answer 1

您可以使用input_file_name()从文件路径中提取S3存储桶名称：

df.withColumn("dates", split(regexp_replace(input_file_name(), "s3://", ""), "/").getItem(0))\
  .show()

我们分割文件名，并获得与存储桶名称相对应的第一部分。

这也可以使用正则表达式s3:\/\/(.+?)\/(.+)完成，第一组是存储桶名称：

df.withColumn("dates", regexp_extract(input_file_name(), "s3:\/\/(.+?)\/(.+)", 1)).show()