使用文件名创建变量 - PySpark答案

【问题标题】：Using filenames to create variable - PySpark使用文件名创建变量 - PySpark
【发布时间】：2020-05-23 13:39:42
【问题描述】：

我有一个文件夹，其中文件被删除（每天，每周），我需要将文件名中格式一致的年份和周/日作为变量添加到我的数据框中。前缀可以更改（例如，sales_report、cash_flow 等），但最后一个字符始终为 YYYY_WW.csv。

例如，对于每周文件，我可以为每个文件手动执行以下操作：

from pyspark.sql.functions import lit

df = spark.read.load('my_folder/sales_report_2019_12.csv', format="csv").withColumn("sales_year", lit(2019)).withColumn("sales_week", lit(12))

我想做相当于使用从文件名右侧开始计数的子字符串函数来解析12 和2019。我是否能够解析这些变量的文件名，然后我可以使用通配符（例如df = spark.read.load('my_folder/sales_report_*.csv', format="csv")）读取文件夹中的所有文件，这将大大简化我的代码。

【问题讨论】：

标签： apache-spark pyspark apache-spark-sql pyspark-dataframes

【解决方案1】：

您可以使用input_file_name() 列和一些字符串函数（如regexp_extract 和substring_index）轻松地从文件名中提取它：

df = spark.read.load('my_folder/*.csv', format="csv")

df = df.withColumn("year_week", regexp_extract(input_file_name(), "\d{4}_\d{1,2}"))\
       .withColumn("sales_year", substring_index(col("year_week"), "_", 1))\
       .withColumn("sales_week", substring_index(col("year_week"), "_", -1))\
       .drop("year_week")

【讨论】：

【解决方案2】：

你可以试试下面的：

import glob
listfiles = glob.glob('my_folder/sales_report_*.csv')
for files in listfiles:
    weekyear = c.split('_',2)[-1].split('_')
    week = weekyear[1]
    year = weekyear[0]
    df = spark.read.load('files', format="csv").withColumn("sales_year", lit(year)).withColumn("sales_week", lit(week))

【讨论】：