PySpark批量读取S3多路径时如何将分区字段作为列加载

逐路径循环读取再union的方案性能瓶颈非常明显：每次循环触发独立读请求，多次union会生成冗余的长血缘执行计划，完全无法利用Spark的批量读IO优化。
Spark原生支持key=value格式的Hive风格分区路径自动识别，不需要手动写逻辑解析路径提取分区值，只要配置正确参数，批量读取时会自动将分区字段作为列加入结果集。

最优方案：开启原生分区发现 #

读取时传入basePath参数指向分区表根目录即可，不管是直接读根路径做分区裁剪，还是传入提前拼接好的指定分区路径列表，Spark都会自动解析路径中的分区键值，零额外开销生成分区列。

写法1：直接读根路径+分区裁剪（无需手动拼路径） #

Spark会自动扫描根路径下的分区元数据，根据过滤条件只读取目标日期范围的数据，比手动拼路径效率更高：

from pyspark.sql.functions import col
from datetime import timedelta

end_date = start_date + timedelta(days=date_count - 1)
rtn_df = spark.read.load(
    s3_prefix,
    basePath=s3_prefix
).filter(
    (col("partition_country") == country)
    & (col("partition_date").between(start_date.strftime("%Y%m%d"), end_date.strftime("%Y%m%d")))
)

写法2：使用预生成的路径列表读取 #

如果你已经提前拼接好了所有待读取的路径列表，只要在load方法中传入basePath参数即可，返回的结果会自动携带partition_country、partition_date两个分区列，不需要手动withColumn添加：

paths = []
for each_day in (start_date + timedelta(d) for d in range(date_count)):
    path = f"{s3_prefix}/partition_country={country}/partition_date={each_day.strftime('%Y%m%d')}"
    paths.append(path)

rtn_df = spark.read.load(
    *paths,
    basePath=s3_prefix
)

兜底方案：手动从文件路径提取分区值 #

如果受场景限制无法配置basePath，可以通过input_file_name()函数获取每条记录对应的S3文件路径，再通过正则提取分区值。该方案会产生额外的路径解析计算开销，仅作为特殊场景下的兜底选项：

from pyspark.sql.functions import input_file_name, regexp_extract

rtn_df = spark.read.load(*paths) \
    .withColumn("partition_date", regexp_extract(input_file_name(), r"partition_date=(\d{8})", 1))

补充：逐路径读取时需要手动添加分区列的原因，就是读取单分区路径时没有指定basePath，Spark无法识别当前路径属于分区结构，自然不会自动生成分区列。

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