如何从PySpark DataFrame列中提取文件扩展名

嘿，这个需求我平时处理得挺多，给你两种实用的实现方式，看哪种更适配你的场景：

方法一：用Python UDF封装os.path.splitext #

因为os.path.splitext是Python标准库的函数，我们可以把它包装成PySpark的UDF来处理DataFrame列：

先导入必要的模块：

import os
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

接着定义专门提取扩展名的UDF：

get_extension = udf(lambda path: os.path.splitext(path)[1], StringType())

最后用withColumn新增一列（比如命名为FileExtension）：

df = df.withColumn("FileExtension", get_extension(df["FullPath"]))

这种方式逻辑直观，如果你已经熟悉os.path.splitext的用法，上手特别快；但要注意，UDF在处理超大规模数据集时，性能可能不如PySpark原生函数。

方法二：使用PySpark内置函数（大数据场景推荐） #

如果你的数据集量级很大，优先用原生函数，避免UDF带来的性能开销。这里推荐两种原生实现：

方式A：正则匹配精准提取 #

用正则表达式匹配路径中最后一个.后面的内容，直接拿到扩展名：

from pyspark.sql.functions import regexp_extract

df = df.withColumn("FileExtension", regexp_extract(df["FullPath"], r'\.([^\.]+)$', 1))

简单解释下正则：\.([^\.]+)$ 会匹配字符串末尾的.，然后捕获后面所有非.的字符，这就是我们要的扩展名。

方式B：分割字符串提取 #

先按.分割路径，再取分割后的最后一段：

from pyspark.sql.functions import split, element_at

df = df.withColumn("FileExtension", element_at(split(df["FullPath"], "\."), -1))

这种方式要注意：如果路径里有多个.（比如doc.v2.pdf），会返回最后一段pdf，和os.path.splitext结果一致；但如果路径本身没有扩展名（比如data/report），会返回整个文件名，你可以根据需求加个判断处理这种情况。

原生函数的性能比UDF好很多，非常适合生产环境的大数据场景。

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