如何使用PySpark将含其他列名的字符串替换为对应列值？

问题场景 #

现有如下PySpark DataFrame：

需要将col1、col2中以%开头的占位符替换为对应列的取值，最终得到仅保留ID、col1、col2的结果表，且避免低效的逐行遍历操作。

高效实现方案 #

PySpark的内置函数支持分布式执行，完全不需要逐行遍历（这种操作会将数据拉到Driver端，严重影响性能，大数据量下甚至会内存溢出）。以下两种方案均基于内置函数实现：

方案1：固定占位符的直接替换

如果需要替换的占位符（如%colA、%colB）数量固定且已知，可以直接链式调用regexp_replace完成替换：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 初始化SparkSession并构建示例DataFrame
spark = SparkSession.builder.appName("placeholder_replace").getOrCreate()
data = [
    ("id_1", "%colA", "t < %colA", "int1", "int3"),
    ("Id_2", "%colB", "t < %colB", "int2", "int4")
]
df = spark.createDataFrame(data, ["ID", "col1", "col2", "colA", "colB"])

# 执行替换并选择目标列
result_df = df.withColumn(
    "col1",
    F.regexp_replace(F.regexp_replace(F.col("col1"), "%colA", F.col("colA")), "%colB", F.col("colB"))
).withColumn(
    "col2",
    F.regexp_replace(F.regexp_replace(F.col("col2"), "%colA", F.col("colA")), "%colB", F.col("colB"))
).select("ID", "col1", "col2")

result_df.show()

执行后输出结果：

+-----+-----+---------+
|   ID| col1|     col2|
+-----+-----+---------+
|id_1 |int1 |t < int1 |
|Id_2 |int4 |t < int4 |
+-----+-----+---------+

方案2：动态适配多占位符

如果需要替换的占位符数量不确定或较多，可以通过动态遍历列名生成替换逻辑，提升代码灵活性：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("dynamic_replace").getOrCreate()
data = [
    ("id_1", "%colA", "t < %colA", "int1", "int3"),
    ("Id_2", "%colB", "t < %colB", "int2", "int4")
]
df = spark.createDataFrame(data, ["ID", "col1", "col2", "colA", "colB"])

# 定义需要处理的目标列和被引用的列（即占位符对应的列）
target_cols = ["col1", "col2"]
ref_cols = [col for col in df.columns if col not in target_cols and col.startswith("col")]

# 动态生成替换逻辑
temp_df = df
for target_col in target_cols:
    for ref_col in ref_cols:
        temp_df = temp_df.withColumn(
            target_col,
            F.regexp_replace(F.col(target_col), f"%{ref_col}", F.col(ref_col))
        )

result_df = temp_df.select("ID", "col1", "col2")
result_df.show()

方案优势 #

两种方案均使用PySpark原生内置函数，完全基于分布式计算框架执行：

• 避免将分布式数据拉取到Driver端处理，不会出现内存溢出问题
• 利用Spark的并行计算能力，处理大数据量时性能远高于逐行遍历

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