在PyArrow中实现按组累计求和（无需转换为Pandas）

嗨，这事儿好办！在PyArrow里不用转成Pandas也能实现和你给出的Pandas代码完全一致的按组累计求和效果，核心是用PyArrow的窗口函数结合分组规则来实现，具体代码如下：

首先，我们先写出完整的实现代码：

import pyarrow as pa
import pyarrow.compute as pc

# 创建你给出的初始PyArrow表
table = pa.table({'a': [1,2,3,4,5,6], 'b': ['x']*3+['y']*3})

# 定义窗口规则：按`b`列分组，同时用行号保证分组内的顺序和原表一致（和Pandas的cumsum行为对齐）
window_spec = pa.window(
    partition_by='b',
    order_by=pc.row_number(table)
)

# 计算按组累计求和
grouped_cumsum = pc.cumsum(table['a'], window=window_spec)

# 可以直接查看结果，和你Pandas的输出完全一样
print(grouped_cumsum.to_pandas())

运行这段代码后，输出结果就是：

0     1
1     3
2     6
3     4
4     9
5    15
Name: a, dtype: int64

关键细节解释 #

• 为什么要加order_by=pc.row_number(table)？因为Pandas的cumsum是按照数据的原始行顺序来累计的，PyArrow的窗口函数需要明确指定排序规则来保证分组内的顺序和原表一致，否则可能出现顺序错乱的情况。
• 如果需要把累计求和的结果添加到原表中，可以用append_column方法：

final_table = table.append_column('a_cumsum', grouped_cumsum)
print(final_table.to_pandas())

输出的完整表会是：

a  b  a_cumsum
0  1  x         1
1  2  x         3
2  3  x         6
3  4  y         4
4  5  y         9
5  6  y        15

备注：内容来源于stack exchange，提问作者ignoring_gravity