如何在PyArrow中实现按组的1索引累积计数（不转Pandas）

数据示例 #

data = {'a': [1,1,2], 'b': [4,5,6]}

需求 #

对列a分组后，生成每组内从1开始的累积计数，效果等价于Pandas中的groupby('a').cumcount() + 1。

Pandas实现参考 #

import pandas as pd

pd.DataFrame(data).groupby('a').cumcount() + 1

PyArrow实现方案 #

无需将初始PyArrow Table转换为Pandas或NumPy数组，可通过PyArrow的计算函数结合NumPy生成序列来实现：

import pyarrow as pa
import numpy as np

# 1. 创建PyArrow Table
table = pa.table(data)

# 2. 获取按'a'分组的分区起止索引
group_indices = pa.compute.partition_indices(table, keys=['a'])

# 3. 为每个分组生成从1开始的计数数组，再合并为一个完整数组
cumcount_array = pa.concat_arrays([
    pa.array(np.arange(1, end - start + 1)) 
    for start, end in group_indices
])

# 4. 将计数列添加到原Table中
result_table = table.append_column('cumcount', cumcount_array)

验证结果：result_table['cumcount']会得到[1, 2, 1]，和Pandas的输出一致。

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