Databricks中如何为CSV加载的表添加文件路径列？

问题背景 #

文件目录结构如下：

FileStore/subfolders/DATE01/filenameA.csv
FileStore/subfolders/DATE01/filenameB.csv
FileStore/subfolders/DATE02/filenameA.csv
FileStore/subfolders/DATE02/filenameB.csv

当前通过以下SQL语句批量加载所有filenameA.csv数据到临时视图：

DROP view IF EXISTS clevertablenameA;
create temporary view clevertablenameA
USING csv
OPTIONS (path "dbfs:/FileStore/subfolders/*/filenameA.csv", header = true)

需要在生成的视图中新增一列，记录每条数据对应的源文件完整路径，预期效果如下：

col1 | col2|....| path
data | data|....| dbfs:/FileStore/subfolders/DATE02/filenameA.csv
data | data|....| dbfs:/FileStore/subfolders/DATE02/filenameA.csv
data | data|....| dbfs:/FileStore/subfolders/DATE02/filenameA.csv

解决方案 #

无需更换数据加载方式，直接使用Databricks内置的input_file_name()函数即可实现需求，提供两种常用写法：

写法1：直接创建带路径列的视图 #

DROP view IF EXISTS clevertablenameA;
CREATE TEMPORARY VIEW clevertablenameA AS
SELECT *, input_file_name() AS path
FROM csv.`dbfs:/FileStore/subfolders/*/filenameA.csv`
OPTIONS (header = true);

写法2：基于基础视图扩展路径列 #

如果需要保留原始无路径的视图，可分两步操作：

-- 1. 创建不含路径列的基础视图
DROP view IF EXISTS clevertablenameA_base;
CREATE TEMPORARY VIEW clevertablenameA_base
USING csv
OPTIONS (path "dbfs:/FileStore/subfolders/*/filenameA.csv", header = true);

-- 2. 基于基础视图生成带路径列的最终视图
DROP view IF EXISTS clevertablenameA;
CREATE TEMPORARY VIEW clevertablenameA AS
SELECT *, input_file_name() AS path
FROM clevertablenameA_base;

补充：PySpark代码实现 #

如果使用PySpark加载数据，也可以通过相同函数实现：

from pyspark.sql.functions import input_file_name

# 加载CSV文件
df = spark.read.csv("dbfs:/FileStore/subfolders/*/filenameA.csv", header=True)
# 添加路径列
df_with_path = df.withColumn("path", input_file_name())
# 创建临时视图
df_with_path.createOrReplaceTempView("clevertablenameA")

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