在Pyspark中的多个列上使用相同的函数重复调用withColumn()

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例DataFrame
data = [("Alice", 25, 160),
        ("Bob", 30, 175),
        ("Charlie", 35, 180)]

df = spark.createDataFrame(data, ["name", "age", "height"])

# 定义要重复调用的函数
def add_suffix(col_name):
    return col(col_name).cast("string").concat(lit("_suffix"))

# 定义要重复调用的列名列表
columns = ["name", "age", "height"]

# 使用循环在多个列上重复调用withColumn()
for column in columns:
    df = df.withColumn(column, add_suffix(column))

# 显示结果
df.show()