apache / apache/arrow

[Python] Efficient way to iterate over groups

Open
#44,676 3 comments 3 reactions 0 assignees View on GitHub
Component: Python Status: needs champion Type: enhancement
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

### Describe the enhancement requested

In pandas / Polars, I can do:
```python
dict(df.group_by(['a', 'b', 'c']).__iter__())
```

There doesn't seem to be a built-in way to do this in PyArrow, hence I'm opening this as a feature request

Concretely, if I have
```python
import pyarrow as pa
tbl = pa.table({'a': [1,1,3], 'b': [4, 4, 4], 'c': [1, 3, 2]})
```
then I'd like a way to end up with
```
{(3,
4): pyarrow.Table
a: int64
b: int64
c: int64
----
a: [[3]]
b: [[4]]
c: [[2]],
(1,
4): pyarrow.Table
a: int64
b: int64
c: int64
----
a: [[1,1]]
b: [[4,4]]
c: [[1,3]]}
```

---

For context, this would be for use in [Narwhals](), where we have tried to come up with a [workaround](https://github.com/narwhals-dev/narwhals/blob/9ecd1767d2cea6956d9ca8b2a12d584705f8023a/narwhals/_arrow/group_by.py#L82-L92), but it does exhibit a noticeable slow-down as the number of grouping keys grows - 3 keys is enough for it to be slower than pandas

### Component(s)

Python

Contributor guide

Open the contributing guide

Research direction

Start with the Narwhals workaround in narwhals/_arrow/group_by.py, lines 82-92, and compare its behavior with the pandas and Polars examples in the issue. Define the PyArrow grouping entry point and verify that it produces tuple-keyed groups as pyarrow.Table objects while avoiding the reported slowdown as grouping keys increase.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.