Add a GROUP BY operator
MEMBER
Hi ! At the moment you can use `to_pandas()` to get a pandas DataFrame that supports `group_by` operations (make sure your dataset fits in memory though)

We use Arrow as a back-end for `datasets` and it doesn't have native group by (see https://github.com/apache/arrow/issues/2189) unfortunately.

I just drafted what it could look like to have `group_by` in `datasets`:
```python
from datasets import concatenate_datasets

def group_by(d, col, join): 
    """from: https://github.com/huggingface/datasets/issues/3644"""
    # Get the indices of each group
    groups = {key: [] for key in d.unique(col)} 
    def create_groups_indices(key, i): 
        groups[key].append(i) 
    d.map(create_groups_indices, with_indices=True, input_columns=col) 
    # Get one dataset object per group
    groups = {key: d.select(indices) for key, indices in groups.items()} 
    # Apply join function
    groups = {
        key: dataset_group.map(join, batched=True, batch_size=len(dataset_group), remove_columns=d.column_names)
        for key, dataset_group in groups.items()
    } 
    # Return concatenation of all the joined groups
    return concatenate_datasets(groups.values())
```

example of usage:
```python

def join(batch): 
    # take the batch of all the examples of a group, and return a batch with one aggregated example
    # (we could aggregate examples into several rows instead of one, if you want)
    return {"total": [batch["i"]]} 

d = Dataset.from_dict({
    "i": [i for i in range(50)],
    "group_key": [i % 4 for i in range(50)],
})
print(group_by(d, "group_key", join))
#                                                total
# 0  [0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48]
# 1  [1, 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45, 49]
# 2     [2, 6, 10, 14, 18, 22, 26, 30, 34, 38, 42, 46]
# 3     [3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 47]
```

Let me know if that helps !

cc @albertvillanova @mariosasko for visibility
2022-02-04T14:56:29Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1030062916
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps49ZYNE
https://api.github.com/repos/huggingface/datasets/issues/comments/1030062916/reactions
2022-02-04T14:57:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1030062916
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
COLLABORATOR
@lhoestq As of PyArrow 7.0.0, `pa.Table` has the [`group_by` method](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.group_by), so we should also consider using that function for grouping. 
2022-02-08T15:06:09Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1032709667
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps49jeYj
https://api.github.com/repos/huggingface/datasets/issues/comments/1032709667/reactions
2022-02-08T15:06:09Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1032709667
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
Any update on this?
2022-06-24T09:56:22Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1165410506
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5FdsDK
https://api.github.com/repos/huggingface/datasets/issues/comments/1165410506/reactions
2022-06-24T09:56:22Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1165410506
https://avatars.githubusercontent.com/u/208336?v=4
https://api.github.com/users/felix-schneider/events{/privacy}
https://api.github.com/users/felix-schneider/followers
https://api.github.com/users/felix-schneider/following{/other_user}
https://api.github.com/users/felix-schneider/gists{/gist_id}
https://github.com/felix-schneider
felix-schneider
MDQ6VXNlcjIwODMzNg==
https://api.github.com/users/felix-schneider/orgs
https://api.github.com/users/felix-schneider/received_events
https://api.github.com/users/felix-schneider/repos
https://api.github.com/users/felix-schneider/starred{/owner}{/repo}
https://api.github.com/users/felix-schneider/subscriptions
User
https://api.github.com/users/felix-schneider
public
MEMBER
You can use https://github.com/mariosasko/datasets_sql by @mariosasko  to go group by operations using SQL queries
2022-06-28T17:39:23Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1169029747
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5Frfpz
https://api.github.com/repos/huggingface/datasets/issues/comments/1169029747/reactions
2022-06-28T17:39:29Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1169029747
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Hi, I have a similar issue as OP but the suggested solutions do not work for my case. Basically, I process documents through a model to extract the last_hidden_state, using the "map" method on a Dataset object, but would like to average the result over a categorical column at the end (i.e. groupby this column).
- A to_pandas() saturates the memory, although it gives me the desired result through a .groupby().apply(np.mean, axis=0) on a smaller use-case,
- The solution posted on Feb 4 is much too slow,
- datasets_sql seems to not like the fact that I'm averaging np.arrays.
So I'm kinda out of "non brute force" options... Any help appreciated
2023-01-18T17:35:11Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1387458980
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5SsvGk
https://api.github.com/repos/huggingface/datasets/issues/comments/1387458980/reactions
2023-01-18T17:36:09Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1387458980
https://avatars.githubusercontent.com/u/11871612?v=4
https://api.github.com/users/jeremylhour/events{/privacy}
https://api.github.com/users/jeremylhour/followers
https://api.github.com/users/jeremylhour/following{/other_user}
https://api.github.com/users/jeremylhour/gists{/gist_id}
https://github.com/jeremylhour
jeremylhour
MDQ6VXNlcjExODcxNjEy
https://api.github.com/users/jeremylhour/orgs
https://api.github.com/users/jeremylhour/received_events
https://api.github.com/users/jeremylhour/repos
https://api.github.com/users/jeremylhour/starred{/owner}{/repo}
https://api.github.com/users/jeremylhour/subscriptions
User
https://api.github.com/users/jeremylhour
public
MEMBER
> Hi, I have a similar issue as OP but the suggested solutions do not work for my case. Basically, I process documents through a model to extract the last_hidden_state, using the "map" method on a Dataset object, but would like to average the result over a categorical column at the end (i.e. groupby this column).
 
If you haven't yet, you could explore using [Polars](https://www.pola.rs/) for this. It's a new DataFrame library written in Rust with Python bindings. It is Pandas like it in many ways ,but does have some biggish differences in syntax/approach so it's definitely not a drop-in replacement. 

Polar's also uses Arrow as a backend but also supports out-of-memory operations; in this case, it's probably easiest to write out your dataset to parquet and then use the polar's `scan_parquet` method (this will lazily read from the parquet file). The thing you get back from that is a `LazyDataFrame` i.e. nothing is loaded into memory until you specify a query and call a  `collect` method. 

Example below of doing a groupby on a dataset which definitely wouldn't fit into memory on my machine:

```
from datasets import load_dataset
import polars as pl

ds = load_dataset("blbooks")
ds['train'].to_parquet("test.parquet")
df = pl.scan_parquet("test.parquet")
df.groupby('date').agg([pl.count()]).collect()
```

>datasets_sql seems to not like the fact that I'm averaging np.arrays.

I am not certain how Polars will handle this either. It does have NumPy support (https://pola-rs.github.io/polars-book/user-guide/howcani/interop/numpy.html) but I assume Polars will need to have at least enough memory in each group you want to average over so you may still end up needing more memory depending on the size of your dataset/groups. 



2023-01-19T09:42:50Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1396693543
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5TP9on
https://api.github.com/repos/huggingface/datasets/issues/comments/1396693543/reactions
2023-01-19T09:42:50Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1396693543
https://avatars.githubusercontent.com/u/8995957?v=4
https://api.github.com/users/davanstrien/events{/privacy}
https://api.github.com/users/davanstrien/followers
https://api.github.com/users/davanstrien/following{/other_user}
https://api.github.com/users/davanstrien/gists{/gist_id}
https://github.com/davanstrien
davanstrien
MDQ6VXNlcjg5OTU5NTc=
https://api.github.com/users/davanstrien/orgs
https://api.github.com/users/davanstrien/received_events
https://api.github.com/users/davanstrien/repos
https://api.github.com/users/davanstrien/starred{/owner}{/repo}
https://api.github.com/users/davanstrien/subscriptions
User
https://api.github.com/users/davanstrien
public
NONE
Hi @davanstrien , thanks a lot, I didn't know about this library and the answer works! I need to try it on the full dataset now, but I'm hopeful. Here's what my code looks like:
```
list_size = 768
df.groupby("date").agg(
    pl.concat_list(
        [
            pl.col("hidden_state")
            .arr.slice(n, 1)
            .arr.first()
            .mean()
            for n in range(0, list_size)
        ]
    ).collect()
```

For some reasons, the following code was giving me a "mean() got unexpected argument 'axis'":
```
df2 = df.groupby('date').agg(
    pl.col("hidden_state").map(np.mean).alias("average_hidden_state")
).collect()

```

EDIT: The solution works on my large dataset, the memory does not crash and the time is reasonable, thanks a lot again!
2023-01-19T12:45:07Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1396920905
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5TQ1JJ
https://api.github.com/repos/huggingface/datasets/issues/comments/1396920905/reactions
2023-01-19T14:10:04Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1396920905
https://avatars.githubusercontent.com/u/11871612?v=4
https://api.github.com/users/jeremylhour/events{/privacy}
https://api.github.com/users/jeremylhour/followers
https://api.github.com/users/jeremylhour/following{/other_user}
https://api.github.com/users/jeremylhour/gists{/gist_id}
https://github.com/jeremylhour
jeremylhour
MDQ6VXNlcjExODcxNjEy
https://api.github.com/users/jeremylhour/orgs
https://api.github.com/users/jeremylhour/received_events
https://api.github.com/users/jeremylhour/repos
https://api.github.com/users/jeremylhour/starred{/owner}{/repo}
https://api.github.com/users/jeremylhour/subscriptions
User
https://api.github.com/users/jeremylhour
public
MEMBER
@jeremylhour glad this worked for you :) 
2023-01-19T14:45:02Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1397087955
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5TRd7T
https://api.github.com/repos/huggingface/datasets/issues/comments/1397087955/reactions
2023-01-19T14:45:02Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1397087955
https://avatars.githubusercontent.com/u/8995957?v=4
https://api.github.com/users/davanstrien/events{/privacy}
https://api.github.com/users/davanstrien/followers
https://api.github.com/users/davanstrien/following{/other_user}
https://api.github.com/users/davanstrien/gists{/gist_id}
https://github.com/davanstrien
davanstrien
MDQ6VXNlcjg5OTU5NTc=
https://api.github.com/users/davanstrien/orgs
https://api.github.com/users/davanstrien/received_events
https://api.github.com/users/davanstrien/repos
https://api.github.com/users/davanstrien/starred{/owner}{/repo}
https://api.github.com/users/davanstrien/subscriptions
User
https://api.github.com/users/davanstrien
public
NONE
I find this functionality missing in my workflow as well and the workarounds with SQL and Polars unsatisfying. Since PyArrow has exposed this functionality, I hope this soon makes it into a release. (:
2023-03-14T14:45:58Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1468237220
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps5Xg4Wk
https://api.github.com/repos/huggingface/datasets/issues/comments/1468237220/reactions
2023-03-14T14:45:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1468237220
https://avatars.githubusercontent.com/u/20798845?v=4
https://api.github.com/users/kuchenrolle/events{/privacy}
https://api.github.com/users/kuchenrolle/followers
https://api.github.com/users/kuchenrolle/following{/other_user}
https://api.github.com/users/kuchenrolle/gists{/gist_id}
https://github.com/kuchenrolle
kuchenrolle
MDQ6VXNlcjIwNzk4ODQ1
https://api.github.com/users/kuchenrolle/orgs
https://api.github.com/users/kuchenrolle/received_events
https://api.github.com/users/kuchenrolle/repos
https://api.github.com/users/kuchenrolle/starred{/owner}{/repo}
https://api.github.com/users/kuchenrolle/subscriptions
User
https://api.github.com/users/kuchenrolle
public
NONE
Any update on this feature? 
2024-03-13T19:34:42Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1995514979
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps528SRj
https://api.github.com/repos/huggingface/datasets/issues/comments/1995514979/reactions
2024-03-13T19:34:42Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1995514979
https://avatars.githubusercontent.com/u/23169815?v=4
https://api.github.com/users/zudi-lin/events{/privacy}
https://api.github.com/users/zudi-lin/followers
https://api.github.com/users/zudi-lin/following{/other_user}
https://api.github.com/users/zudi-lin/gists{/gist_id}
https://github.com/zudi-lin
zudi-lin
MDQ6VXNlcjIzMTY5ODE1
https://api.github.com/users/zudi-lin/orgs
https://api.github.com/users/zudi-lin/received_events
https://api.github.com/users/zudi-lin/repos
https://api.github.com/users/zudi-lin/starred{/owner}{/repo}
https://api.github.com/users/zudi-lin/subscriptions
User
https://api.github.com/users/zudi-lin
public
MEMBER
We added a proper Polars integration at #3334 if it can help:
```python
>>> from datasets import load_dataset
>>> ds = load_dataset("TheBritishLibrary/blbooks", "1700_1799", split="train")
>>> ds.to_polars().groupby('date').len()
┌─────────────────────┬──────┐
│ date                ┆ len  │
│ ---                 ┆ ---  │
│ datetime[ms]        ┆ u32  │
╞═════════════════════╪══════╡
│ 1796-01-01 00:00:00 ┆ 5831 │
│ 1775-01-01 00:00:00 ┆ 4697 │
│ 1749-01-01 00:00:00 ┆ 1118 │
│ 1740-01-01 00:00:00 ┆ 713  │
│ 1714-01-01 00:00:00 ┆ 865  │
│ …                   ┆ …    │
│ 1795-01-01 00:00:00 ┆ 5930 │
│ 1754-01-01 00:00:00 ┆ 1373 │
│ 1780-01-01 00:00:00 ┆ 1970 │
│ 1734-01-01 00:00:00 ┆ 1047 │
│ 1719-01-01 00:00:00 ┆ 1235 │
└─────────────────────┴──────┘
```

2024-03-14T13:39:21Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-1997484590
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps53DzIu
https://api.github.com/repos/huggingface/datasets/issues/comments/1997484590/reactions
2024-03-14T13:39:21Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1997484590
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
Umm... did any responses GET REQUESTS? I cannot understand why 'integrations' are mentioned.
2025-01-02T01:06:50Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-2567207825
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps6ZBHuR
https://api.github.com/repos/huggingface/datasets/issues/comments/2567207825/reactions
2025-01-02T01:06:50Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2567207825
https://avatars.githubusercontent.com/u/65698076?v=4
https://api.github.com/users/sorryhyun/events{/privacy}
https://api.github.com/users/sorryhyun/followers
https://api.github.com/users/sorryhyun/following{/other_user}
https://api.github.com/users/sorryhyun/gists{/gist_id}
https://github.com/sorryhyun
sorryhyun
MDQ6VXNlcjY1Njk4MDc2
https://api.github.com/users/sorryhyun/orgs
https://api.github.com/users/sorryhyun/received_events
https://api.github.com/users/sorryhyun/repos
https://api.github.com/users/sorryhyun/starred{/owner}{/repo}
https://api.github.com/users/sorryhyun/subscriptions
User
https://api.github.com/users/sorryhyun
public
NONE
@lhoestq so does `to_polars` work with memory mapping? Because to_pandas doesn't, does it?
2025-01-27T22:45:47Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-2617056085
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps6b_RtV
https://api.github.com/repos/huggingface/datasets/issues/comments/2617056085/reactions
2025-01-28T00:02:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2617056085
https://avatars.githubusercontent.com/u/3231217?v=4
https://api.github.com/users/JulesGM/events{/privacy}
https://api.github.com/users/JulesGM/followers
https://api.github.com/users/JulesGM/following{/other_user}
https://api.github.com/users/JulesGM/gists{/gist_id}
https://github.com/JulesGM
JulesGM
MDQ6VXNlcjMyMzEyMTc=
https://api.github.com/users/JulesGM/orgs
https://api.github.com/users/JulesGM/received_events
https://api.github.com/users/JulesGM/repos
https://api.github.com/users/JulesGM/starred{/owner}{/repo}
https://api.github.com/users/JulesGM/subscriptions
User
https://api.github.com/users/JulesGM
public
MEMBER
According to the [polars docs](https://docs.pola.rs/api/python/dev/reference/api/polars.from_arrow.html):

> This operation will be zero copy for the most part. Types that are not supported by Polars may be cast to the closest supported type.

which means that for the most part the memory mapped data is not copied, so yes it works with memory mapping :)
2025-01-28T11:39:48Z
https://github.com/huggingface/datasets/issues/3644#issuecomment-2618759296
https://api.github.com/repos/huggingface/datasets/issues/3644
IC_kwDODunzps6cFxiA
https://api.github.com/repos/huggingface/datasets/issues/comments/2618759296/reactions
2025-01-28T11:39:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2618759296
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
**Is your feature request related to a problem? Please describe.**
Using batch mapping, we can easily split examples. However, we lack an appropriate option for merging them back together by some key. Consider this example:

```python
# features:
# {
#    "example_id": datasets.Value("int32"),
#    "text": datasets.Value("string")
# }

ds = datasets.Dataset()


def split(examples):
    sentences = [text.split(".") for text in examples["text"]]
    return {
        "example_id": [
            example_id
            for example_id, sents in zip(examples["example_id"], sentences)
            for _ in sents
        ],
        "sentence": [sent for sents in sentences for sent in sents],
        "sentence_id": [i for sents in sentences for i in range(len(sents))],
    }


split_ds = ds.map(split, batched=True)


def process(examples):
    outputs = some_neural_network_that_works_on_sentences(examples["sentence"])
    return {"outputs": outputs}


split_ds = split_ds.map(process, batched=True)
```

I have a dataset consisting of texts that I would like to process sentence by sentence in a batched way. Afterwards, I would like to put it back together as it was, merging the outputs together.

**Describe the solution you'd like**
Ideally, it would look something like this:

```python
def join(examples):
    order = np.argsort(examples["sentence_id"])
    text = ".".join(examples["text"][i] for i in order)
    outputs = [examples["outputs"][i] for i in order]
    return {"text": text, "outputs": outputs}


ds = split_ds.group_by("example_id", join)
```

**Describe alternatives you've considered**
Right now, we can do this:
```python
def merge(example):
    meeting_id = example["example_id"]
    parts = split_ds.filter(lambda x: x["example_id"] == meeting_id).sort("segment_no")
    return {"outputs": list(parts["outputs"])}

ds = ds.map(merge)
```

Of course, we could process the dataset like this:

```python
def process(example):
    outputs = some_neural_network_that_works_on_sentences(example["text"].split("."))
    return {"outputs": outputs}

ds = ds.map(process, batched=True)
```

However, that does not allow using an arbitrary batch size and may lead to very inefficient use of resources if the batch size is much larger than the number of sentences in one example.

I would very much appreciate some kind of group by operator to merge examples based on the value of one column.
