如何高效迭代PyArrow表的分组？对比Pandas groupby实现

在Pandas里用groupby做分组迭代很方便，但你目前用PyArrow时遍历唯一值再过滤的方式，确实会多次扫描表，效率不高。这里有两种更高效的实现方式：

方法一：利用PyArrow原生GroupBy对象获取分组索引 #

PyArrow的Table.group_by()可以直接生成分组对象，通过其groups属性能一次性获取所有分组的行索引，之后用take()提取子表，只需要扫描一次分组键：

import pyarrow as pa
import pyarrow.compute as pc

data = {'a': [1, 1, 1, 2, 2], 'b': [2, 4, 3, 5, 6]}
tbl = pa.table(data)

# 创建分组对象
grouped = tbl.group_by('a')
# 获取分组值与对应行索引的字典
group_indices = grouped.groups

# 迭代每个分组
for group_key, indices in group_indices.items():
    sub_tbl = tbl.take(indices)
    print(sub_tbl)

输出结果和你之前的一致，但只需要扫描一次a列来计算索引，后续的take操作是基于索引的高效切片，无需重复扫表。

方法二：排序后拆分连续分组（适合超大数据集） #

如果你的数据量极大，可以先按分组键排序，让相同分组值的行连续排列，之后只需一次扫描就能拆分所有分组，避免多次过滤：

import pyarrow as pa
import pyarrow.compute as pc

data = {'a': [1, 1, 1, 2, 2], 'b': [2, 4, 3, 5, 6]}
tbl = pa.table(data)

# 按分组键排序
sorted_tbl = tbl.sort_by('a')
a_col = sorted_tbl['a']

# 计算分组边界：找出相邻元素不同的位置
boundary_mask = pc.not_equal(a_col, pc.shift(a_col, 1))
boundaries = pc.index(boundary_mask).to_pylist()

# 生成拆分的起始/结束索引
split_positions = [0] + boundaries + [len(sorted_tbl)]

# 迭代拆分每个连续分组
for i in range(len(split_positions) - 1):
    start = split_positions[i]
    end = split_positions[i+1]
    sub_tbl = sorted_tbl.slice(start, end - start)
    print(sub_tbl)

这种方法的时间复杂度主要来自排序（O(N log N)），但后续拆分仅需一次线性扫描，比多次过滤的O(N*K)（K为分组数）高效得多，适合处理大规模数据。

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