如何在PyArrow中按列分组后选择每组前K个最大值（不使用Pandas）

你可以通过以下纯PyArrow原生API的步骤实现需求：

实现步骤 #

1. 分组聚合生成列表：按id分组，将每组的score聚合为列表（这部分你已经完成）。
2. 排序并截取Top K：对每个分组的score列表降序排序，再截取前K个元素。
3. 展开列表并关联分组键：将处理后的嵌套列表展开，同时匹配对应的分组id，重构为目标表格。

完整代码示例 #

import pyarrow as pa
import pyarrow.compute as pc

# 构造原始数据集
df = pa.table([
    pa.array([1, 1, 1, 1, 1, 2, 2, 2, 2, 2]), 
    pa.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
], names=["id", "score"])

K = 2

# 1. 分组聚合得到每个id对应的score列表
grouped = df.group_by("id").aggregate([("score", "list")])

# 2. 对每个列表降序排序，截取前K个元素
sorted_top_k = pc.list_slice(
    pc.list_sort(grouped["score_list"], sort_order="descending"),
    start=0, length=K
)

# 3. 展开列表并关联对应id
flattened_scores = pc.list_flatten(sorted_top_k)
parent_indices = pc.list_parent_indices(sorted_top_k)
result_ids = pc.take(grouped["id"], parent_indices)

# 构造最终结果表
result = pa.table({"id": result_ids, "score": flattened_scores})

# 查看结果
print(result.to_pandas())

代码说明 #

• pc.list_sort(..., sort_order="descending")：对每个分组的score列表做降序排序，确保最大值排在前面。
• pc.list_slice()：截取每个列表的前K个元素，保留每组的Top K最大值。
• pc.list_flatten()：把嵌套的列表展开为一维数组，得到所有符合要求的score值。
• pc.list_parent_indices()：获取展开后每个元素对应的原始分组索引，再通过pc.take()匹配对应的id，保证id和score一一对应。

运行后会输出期望结果：

id  score
0   1      4
1   1      3
2   2      9
3   2      8

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