如何在PySpark中将DataFrame单行转换为带新列的多行数据

关于逐行处理的合理性说明 #

• 不可取的是将全量数据通过collect()算子拉取到Driver端做本地遍历，超大规模数据集下会直接触发Driver内存溢出，完全不符合Spark分布式计算的设计逻辑
• 分布式层面的行级处理是Spark的标准操作，你需求的UDTF能力属于典型的分布式行转换算子，全程在集群节点并行执行，不会拉取数据到本地，完全适配超大流量场景

推荐实现方案 #

方案1：PySpark 3.0+ 原生UDTF（最优解，完全对标Snowflake UDTF） #

PySpark 3.0及以上版本原生支持用户自定义表函数（UDTF），支持输入多列、输出多行多列，经过Catalyst引擎优化，性能最优。
示例代码：

from pyspark.sql.functions import udtf, col
from pyspark.sql.types import IntegerType, StringType

# 定义UDTF，指定返回的多列结构
@udtf(returnType="new_col1: int, new_col2: string")
class CustomUDTF:
    def eval(self, input_col1: int, input_col2: str):
        # 此处编写你的自定义处理逻辑，支持根据输入参数生成任意行数的返回值
        yield (input_col1 * 1, f"{input_col2}_suffix1")
        yield (input_col1 * 2, f"{input_col2}_suffix2")
        yield (input_col1 * 3, f"{input_col2}_suffix3")

# 调用UDTF，自动将单行拆分为多行，原有列完整保留
result_df = df.join(CustomUDTF.explode(col("input_col1"), col("input_col2")))

方案2：UDF返回结构化数组 + explode（兼容PySpark 3.0以下版本） #

如果使用低版本PySpark，可以通过UDF返回结构化数组，再配合explode拆分为多行实现相同效果，性能略低于原生UDTF但仍适配大规模数据集，也解决了原生explode功能过于单一的问题。
示例代码：

from pyspark.sql.functions import udf, explode, col
from pyspark.sql.types import ArrayType, StructType, StructField, IntegerType, StringType

# 定义UDF返回的数组结构，每个元素对应生成的一行数据
output_schema = ArrayType(StructType([
    StructField("new_col1", IntegerType(), nullable=False),
    StructField("new_col2", StringType(), nullable=False)
]))

@udf(returnType=output_schema)
def custom_udf(input_col1: int, input_col2: str):
    # 自定义处理逻辑，返回多组结果的列表
    return [
        (input_col1 * 1, f"{input_col2}_suffix1"),
        (input_col1 * 2, f"{input_col2}_suffix2"),
        (input_col1 * 3, f"{input_col2}_suffix3")
    ]

# 执行转换：生成数组 -> 拆分为多行 -> 展开新字段
result_df = df.withColumn("temp_output", explode(custom_udf(col("input_col1"), col("input_col2")))) \
              .select("*", col("temp_output.new_col1"), col("temp_output.new_col2")) \
              .drop("temp_output")

避坑提示 #

• 不要使用forEach算子：它属于无返回值的Action算子，仅用于写出数据到外部系统，不适合做DataFrame转换
• transform方法是面向整个DataFrame的批量转换接口，不用于行级处理，不需要强行适配该方法实现需求
• 无特殊需求不要使用RDD算子：DataFrame原生的UDTF/UDF已经过引擎优化，性能比RDD的map类算子高30%以上
• 自定义逻辑中的重初始化操作（如加载模型、读取配置）建议在UDTF的__init__方法或者mapPartitions算子中实现，避免每行重复初始化导致性能损耗

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