在PySpark中用自定义函数/UDF处理同RequestID分组并保留原行数

现有数据集示例：

RequestID  |  RankedItem    | ... more columns ... |
-----------------------------------------------------
RequestID 1  |  RankedItem 11 |
RequestID 1  |  RankedItem 12 |
RequestID 1  |  RankedItem 13 |
RequestID 2  |  RankedItem 21 |
RequestID 2  |  RankedItem 22 |
RequestID 2  |  RankedItem 23 |

需求：对同一RequestID分组内的所有行应用自定义函数/UDF，利用该分组下的全部信息为每一行生成新列，最终输出行数和原数据集完全一致，示例输出如下：

RequestID  |  RankedItem    | ... more columns ... | *new column*
-------------------------------------------------------------
RequestID 1  |  RankedItem 11 |                      | some_value
RequestID 1  |  RankedItem 12 |                      | some_other_value
RequestID 1  |  RankedItem 13 |                      | another_value

已知apply是逐行操作，groupby会对数据分组但不需要保留分组形式，如何实现上述需求？

以Pandas为例，可通过以下两种方式实现：

方法1：groupby + apply（灵活处理复杂逻辑） #

利用groupby按RequestID分组后，让自定义函数接收整个分组的DataFrame进行处理，返回包含新列的分组DataFrame，Pandas会自动将所有分组结果拼接回原行数的完整DataFrame：

1. 定义处理分组的自定义函数：

def process_group(group_df):
    # 示例逻辑：基于分组内所有RankedItem，生成对应自定义值
    # 实际可替换为任意需要用到全组数据的逻辑
    item_nums = group_df['RankedItem'].str.extract('(\d+)').astype(int)
    group_df['new_column'] = item_nums.apply(lambda x: f"custom_value_{x}")
    return group_df

2. 分组应用函数：

# group_keys=False 避免生成冗余的分组索引列
result_df = df.groupby('RequestID', group_keys=False).apply(process_group)

方法2：groupby + transform（适合简单分组计算） #

如果自定义逻辑仅需返回与分组行数一致的序列，可使用transform方法，它会自动将分组级别的计算结果映射到每一行，无需手动拼接：

# 示例：计算每个RequestID分组内RankedItem的最大值，映射到组内每一行
df['new_column'] = df.groupby('RequestID')['RankedItem'].transform(
    lambda x: max(x.str.extract('(\d+)').astype(int))
)

transform的优势是代码更简洁，但仅适用于输出单一序列的场景；若需要多列交互处理或生成多个新列，优先选择groupby + apply。

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