如何将PySpark DataFrame数组元素映射到列并格式化字符串

你可以通过两种方式实现这个需求：自定义UDF（直观易理解）或者纯Spark内置函数（性能更优，推荐生产环境使用）。

方法一：自定义UDF #

通过定义UDF动态提取参数对应的列值，再格式化SQL语句：

from pyspark.sql import SparkSession, functions as F, Row, types as T

# 初始化SparkSession（未初始化时执行）
spark = SparkSession.builder.appName("FormatStatement").getOrCreate()

# 构建原始DataFrame
sdf1 = spark.sparkContext.parallelize([
    ["toto", "tata", ["table", "column"], "SELECT {1} FROM {0}"],
    ["titi", "tutu", ["table", "column"], "SELECT {1} FROM {0}"]
]).toDF(["table", "column", "parameters", "statement"])

# 定义语句格式化函数
def format_statement(params, stmt, row):
    # 根据参数列表提取对应列的值
    values = [row[col] for col in params]
    # 填充模板生成最终语句
    return stmt.format(*values)

# 注册UDF
format_udf = F.udf(format_statement, T.StringType())

# 生成结果列
sdf2 = sdf1.withColumn(
    "result",
    format_udf(F.col("parameters"), F.col("statement"), F.struct(F.col("*")))
)

# 查看输出结果
sdf2.show(truncate=False)

方法二：纯Spark内置函数 #

利用Spark原生函数实现，避免UDF的序列化开销：

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.appName("FormatStatement").getOrCreate()

# 构建原始DataFrame
sdf1 = spark.sparkContext.parallelize([
    ["toto", "tata", ["table", "column"], "SELECT {1} FROM {0}"],
    ["titi", "tutu", ["table", "column"], "SELECT {1} FROM {0}"]
]).toDF(["table", "column", "parameters", "statement"])

# 生成结果列
sdf2 = sdf1.withColumn(
    "result",
    F.expr("""
        format_string(
            statement,
            transform(parameters, col_name -> element_at(named_struct(*array(columns)), col_name))
        )
    """)
)

# 查看输出结果
sdf2.show(truncate=False)

两种方法都会输出你需要的结果：

+-----+------+---------------+-------------------+---------------------+
|table|column|parameters     |statement          |result               |
+-----+------+---------------+-------------------+---------------------+
|toto |tata  |[table, column]|SELECT {1} FROM {0}|SELECT tata FROM toto|
|titi |tutu  |[table, column]|SELECT {1} FROM {0}|SELECT tutu FROM titi|
+-----+------+---------------+-------------------+---------------------+

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