PySpark按条件透视DataFrame：保留每组同类型最多3个值

咱们一步步来实现你想要的结果，核心是先过滤无效数据、分组排序取前3条，再进行透视操作：

1. 过滤无效的type值 #

首先处理new_df里空的type（包括空字符串和全空格的情况），这些不需要出现在最终结果里，直接过滤掉：

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

# 过滤空/空白的type
filtered_df = new_df.filter(F.trim(F.col("type")) != "")

2. 给每个id+type组添加序号并保留前3条 #

用窗口函数给每个id+type分组内的记录按s_id排序（你可以根据实际需求调整排序字段），生成行号后只保留前3条：

# 定义窗口规则：按id和type分组，按s_id排序
window_spec = Window.partitionBy("id", "type").orderBy("s_id")

# 添加行号，过滤出前3条记录
ranked_df = filtered_df.withColumn("row_num", F.row_number().over(window_spec)) \
                       .filter(F.col("row_num") <= 3)

3. 构造透视列名 #

把type和行号拼接成最终的列名格式（比如android1、ios2）：

pivot_col_df = ranked_df.withColumn("pivot_column", F.concat(F.col("type"), F.col("row_num")))

4. 执行透视操作 #

按id分组，用构造好的列名做透视，聚合取s_id的值（用first是因为每个分组位置只有一个值，用max/last也可以）：

# 生成透视表
pivoted_df = pivot_col_df.groupBy("id") \
                         .pivot("pivot_column") \
                         .agg(F.first("s_id"))

5. 调整列顺序（可选） #

如果需要和你期望的列顺序完全一致，可以手动指定列的排列：

# 定义目标列顺序
target_columns = ["id", "android1", "android2", "android3", "ios1", "ios2", "ios3"]
final_df = pivoted_df.select(*target_columns)

最终结果 #

执行final_df.show()就能得到你想要的输出：

+---+--------+--------+--------+----+----+----+
|id |android1|android2|android3|ios1|ios2|ios3|
+---+--------+--------+--------+----+----+----+
|1 |15      | null   | null   | 11 | 12 |null|
|2 |18      | null   | null   | 21 |null|null|
+---+--------+--------+--------+----+----+----+

关键说明 #

• 过滤空type是为了避免生成无意义的空字符串列，符合你的预期；
• 窗口函数的row_number()确保每个id+type组最多保留3条记录；
• 拼接type和行号作为透视列，正好匹配你需要的列命名规则。

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