基于正则表达式列表批量重命名PySpark DataFrame列

需求说明 #

现有PySpark DataFrame，需基于指定正则表达式列表批量重命名列，将所有匹配模式的列统一命名为price_<捕获组内容>格式，非匹配列（如item_name）保留原名。

给定正则表达式模式 #

patterns = [
    "price-usd-([0-9]+)", 
    "list_price_([0-9]+)", 
    "price_per_([0-9]+)_units", 
    "pricefor([0-9]+)", 
    "([0-9]+)_plus_price", 
    "break_price_([0-9]+)", 
    "price_break_pricing_([a-z]+)"
]

原DataFrame结构与数据 #

期望输出DataFrame结构与数据 #

解决方案代码 #

from pyspark.sql import SparkSession
import re

# 初始化SparkSession
spark = SparkSession.builder.appName("BatchRenameColumns").getOrCreate()

# 构造原DataFrame（实际场景中可直接使用已有DataFrame）
data = [
    ("Samsung Z", 10000, 5, 9000, 10, 7000, 4, 21),
    ("Moto G4", 12000, 10, 10000, 20, 6000, 3, 43),
    ("Mi 4i", 15000, 8, 12000, 20, 10000, 5, 25),
    ("Moto G3", 20000, 5, 18000, 12, 15000, 10, 15)
]
original_columns = [
    "item_name", "price-usd-1", "break_price_7", "pricefor5", 
    "price_per_9_units", "price_break_pricing_a", "2_plus_price", "list_price_8"
]
df = spark.createDataFrame(data, schema=original_columns)

# 定义正则模式列表
patterns = [
    "price-usd-([0-9]+)", 
    "list_price_([0-9]+)", 
    "price_per_([0-9]+)_units", 
    "pricefor([0-9]+)", 
    "([0-9]+)_plus_price", 
    "break_price_([0-9]+)", 
    "price_break_pricing_([a-z]+)"
]

# 定义列名转换函数
def transform_col_name(col_name):
    for pattern in patterns:
        match_result = re.match(pattern, col_name)
        if match_result:
            return f"price_{match_result.group(1)}"
    return col_name

# 批量重命名列
renamed_df = df.select(*[df[col].alias(transform_col_name(col)) for col in df.columns])

# 打印结果
renamed_df.show()

代码说明 #

1. 正则匹配逻辑：遍历每个列名，逐个匹配给定的正则模式，一旦匹配成功则提取模式中的捕获组内容，拼接成price_xxx格式的新列名。
2. 非匹配列处理：如果列名不匹配任何正则模式（如item_name），则保留原列名。
3. 批量重命名：通过select方法结合列表推导式，一次性完成所有列的重命名操作，效率优于循环调用withColumnRenamed。

内容的提问来源于stack exchange，提问作者Ramineni Ravi Teja