迭代 PyArrow 表组

Question

在 pandas 中，我可以迭代

groupby

中的组：

In [3]: import pandas as pd

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

In [5]: df = pd.DataFrame(data)

In [6]: for _, sub_df in df.groupby('a'):
   ...:     print(sub_df)
   ...:
   a  b
0  1  2
1  1  4
2  1  3
   a  b
3  2  5
4  2  6

PyArrow 有没有有效的方法来做到这一点？说我从

开始

tbl = pa.table(data)

我能带的就是：

In [16]: for x in pc.unique(tbl['a']):
    ...:     print(tbl.filter(pc.equal(tbl['a'], x)))
    ...:
pyarrow.Table
a: int64
b: int64
----
a: [[1,1,1]]
b: [[2,4,3]]
pyarrow.Table
a: int64
b: int64
----
a: [[2,2]]
b: [[5,6]]

但这涉及多次扫描整个

'a'

列......有没有更高效的方法？

Answer 1

Table.group_by可以聚合成列表数组。

>>> data = {'a': [1, 1, 1, 2, 2], 'b': [2, 4, 3, 5, 6]}
>>> groups = pa.table(data).group_by(['a']).aggregate([('b', 'list')])
>>> groups
pyarrow.Table
a: int64
b_list: list<item: int64>
  child 0, item: int64
----
a: [[1,2]]
b_list: [[[2,4,3],[5,6]]]

组位于单个表中，但行可以迭代。

>>> for row in groups.to_struct_array():
...     print(row['a'], row['b_list'])
...     
1 [2, 4, 3]
2 [5, 6]