PySpark实现多值ID行转多列扁平化（低内存方案）

核心思路 #

通过窗口函数为每个ID下的重复行生成唯一序号，再利用pivot操作转置列，同时明确指定需要生成的列数，避免Spark扫描所有可能的序号值，从而控制内存占用。以下是具体步骤：

1. 为每个ID内的行添加序号 #

使用row_number()窗口函数，按ID分组后为每行生成从1开始的序号，对应最终的Text_1、Text_2等列：

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

# 假设原始DataFrame名为df
window_spec = Window.partitionBy("ID").orderBy("Text")  # 可根据实际需求调整排序规则
df_with_rank = df.withColumn("text_idx", F.row_number().over(window_spec))

2. 利用Pivot转置列（指定最大列数） #

你已统计出需要扩展的最大列数（示例中为3），在pivot时明确传入列序号列表，避免Spark自动扫描所有可能的序号值，大幅减少内存开销：

# 假设你统计得到的最大列数为max_cols
max_cols = 3

pivoted_df = df_with_rank.groupBy("ID").pivot(
    "text_idx",
    values=[str(i) for i in range(1, max_cols + 1)]  # 明确指定要生成的列序号
).agg(F.first("Text"))  # 每个ID+序号组合仅一行，用first()/max()/min()均可

3. 重命名列名 #

将默认的数字列名改为Text_1、Text_2格式：

final_df = pivoted_df.select(
    "ID",
    *[F.col(str(i)).alias(f"Text_{i}") for i in range(1, max_cols + 1)]
)

4. 自动处理空值 #

对于没有对应序号行的ID（如示例中的ID=3），pivot会自动填充null，无需额外处理即可符合预期结果。

完整代码示例 #

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

# 构造原始DataFrame
data = [
    (1, "some text"),
    (1, "more text"),
    (1, "still more text"),
    (2, "some text"),
    (2, "still more text"),
    (3, None)
]
df = spark.createDataFrame(data, ["ID", "Text"])

# 统计最大列数（若已提前计算可跳过此步）
max_cols = df.groupBy("ID").agg(F.count("Text").alias("cnt")).agg(F.max("cnt")).collect()[0][0]

# 添加行序号
window_spec = Window.partitionBy("ID").orderBy("Text")
df_with_rank = df.withColumn("text_idx", F.row_number().over(window_spec))

# Pivot转置
pivoted_df = df_with_rank.groupBy("ID").pivot(
    "text_idx",
    values=[str(i) for i in range(1, max_cols + 1)]
).agg(F.first("Text"))

# 重命名列并输出结果
final_df = pivoted_df.select(
    "ID",
    *[F.col(str(i)).alias(f"Text_{i}") for i in range(1, max_cols + 1)]
)

final_df.show()

执行后输出结果与你期望的完全一致，且所有操作均为分布式执行，不会将全量数据加载到单节点内存，适合大数据场景。

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