Add reduce function
State: closed

This PR closes #5496 .

I tried to imitate the `reduce`-method from `functools`, i.e. the function input must be a binary operation. I assume that the input type has an empty element, i.e. `input_type()` is defined, as the acumulant is instantiated as this object - im not sure that is this a reasonable assumption?

If `batched= True`  the reduction of each shard is _not_ returned, but the reduction of the entire dataset. I was unsure wether this was an intuitive API, or it would make more sense to return the reduction of each shard?

Comments:
I agree that it would be a good idea to introduce a `combiner` argument in another PR.

I did take quite a lot of inspiration from the implementation of  `map`, but it did not seem obvious how to resuse `map` for the implementation. Do you have any suggestions, i could give a try?

Those were exactly my thoughts, regarding the non-obvious initializer for batched and formatted datasets, so i agree! I'll  introduce a `initializer` argument, and have it mandatory when `batched=True`.

I added `initializer`. It is optional for `batched=False` and mandatory for `batched=True`. It has to be of the same length as `input_columns`, if `input_columns=None` it has to have the same length as `_data.column_names`. 

If the initializer is not set for `batched=False` the first example is set as the `initializer`. 

The initializer is used to initiliaze for each shard, so that means if that:
```python
dset = Dataset.from_dict({"x": [1, 2, 3]})
sum_reduce = lambda x, y: x + y
reduction = dset.reduce(sum_reduce, batched=True, initializer=1, input_columns='x', num_proc=2)
# reduction is 8, i.e. reduction + num_proc * initializer
```

> I added initializer. It is optional for batched=False and mandatory for batched=True. It has to be of the same length as input_columns, if input_columns=None it has to have the same length as _data.column_names.
> 
> If the initializer is not set for batched=False the first example is set as the initializer.

Sounds good to me !

> The initializer is used to initiliaze for each shard, so that means if that:
> 
> ```python
> dset = Dataset.from_dict({"x": [1, 2, 3]})
> sum_reduce = lambda x, y: x + y
> reduction = dset.reduce(sum_reduce, batched=True, initializer=1, input_columns='x', num_proc=2)
> # reduction is 8, i.e. reduction + num_proc * initializer
> ```

Hmm this can be confusing for some users. Maybe we should consider making `combiner` mandatory for multiprocessing.

If we agree on this, maybe for this PR you can either:
- remove multiprocessing (and we add combiner + multiprocessing in a subsequent PR)
- OR add `combiner` directly

Maybe we can get more feedback from @huggingface/datasets as well

> > I added initializer. It is optional for batched=False and mandatory for batched=True. It has to be of the same length as input_columns, if input_columns=None it has to have the same length as _data.column_names.
> > If the initializer is not set for batched=False the first example is set as the initializer.
> 
> Sounds good to me !
> 
> > The initializer is used to initiliaze for each shard, so that means if that:
> > ```python
> > dset = Dataset.from_dict({"x": [1, 2, 3]})
> > sum_reduce = lambda x, y: x + y
> > reduction = dset.reduce(sum_reduce, batched=True, initializer=1, input_columns='x', num_proc=2)
> > # reduction is 8, i.e. reduction + num_proc * initializer
> > ```
> 
> Hmm this can be confusing for some users. Maybe we should consider making `combiner` mandatory for multiprocessing.
> 
> If we agree on this, maybe for this PR you can either:
> 
> * remove multiprocessing (and we add combiner + multiprocessing in a subsequent PR)
> * OR add `combiner` directly
> 
> Maybe we can get more feedback from @huggingface/datasets as well

I think i prefer adding `combiner` in this PR. I think ill make `combiner` mandatory for `batched=True`, instead of assuming that `combiner=function`. Ill look at this one of the coming days. Also at some point i have to define `reduce` for `DatasetDict`, and not just `Dataset`.

I added the `combiner` parameter as described. I added some examples in the docstring, as i felt it might still be a bit confusing what happens during multiprocessing / batching.

Still need to look at `DatasetDict`.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_5533). All of your documentation changes will be reflected on that endpoint.

Feel free to merge `main` into your branch - we fixed some CI failures today

The proposed API doesn't seem intuitive to me - one can already use `functools.reduce` or `Dataset.map` for this purpose ([Colab](https://colab.research.google.com/drive/1jCLv31Y4cDfqD0lhO0AnqEv3Or-LLvWe?usp=sharing) with examples), so perhaps we could have a section in the docs that uses these methods to perform reductions rather than introducing a new method (which needs to be maintained later)

Thanks for sharing this google colab, it has nice examples !

Though I still think `functools.reduce` with multiprocessing can be a pain - we offer something easier here:
- no need to use a pool yourself
- no need to use `map` just to iterate on the dataset (not its main purpose)
- native support for lambdas (using dill)
- the combiner is **mandatory** for multiprocessing to avoid ending up with an incorrect result as in your example

However I agree that maintaining this can be challenging, especially if you think about how `map` already is, and if we also have to deal with dataset formatting.

> native support for lambdas (using dill)

Replacing `multiprocessing` with `multiprocess` in the example would allow that.

> no need to use map just to iterate on the dataset (not its main purpose)

Not the main purpose, but this was mentioned as a "feature" in the previous docs if I remember.

And all this is related to the multi-processing case, which we can document.

Besides the linked issue, I can't find requests for `Dataset.reduce`, which makes me think `functools.reduce` does the job for most users.

> Besides the linked issue, I can't find requests for Dataset.reduce, which makes me think functools.reduce does the job for most users.

I think @srush was looking for a way to do a word count but ended up using a single processed `map`. I also saw some users on the forum wanting to compute `max`

> Not the main purpose, but this was mentioned as a "feature" in the previous docs if I remember.
> 
> And all this is related to the multi-processing case, which we can document.

Yup indeed

While counting is one example, I often find I want to compute different statistics over a dataset. This seems like a natural way to do it in a stateless manner.


I guess you could use functools reduce, but that wouldn't allow batching, right?

I've updated the [Colab](https://colab.research.google.com/drive/1jCLv31Y4cDfqD0lhO0AnqEv3Or-LLvWe?usp=sharing) with an example that reduces batches with `map` and then computes the final result. It would be nice to have a similar example (explained in detail) in the docs to show the full power of `map`.

Plus, for simple reductions such as `max`, one can do `pc.max(ds.with_format("arrow")["col"])` to directly get the result (without loading the entire column in RAM).

@srush 

> I guess you could use functools reduce, but that wouldn't allow batching, right?

You can use `.iter(batch_size)` to get batches
 

That `functools` tools example is clean. I didn't know about `iter`. That would handle my use case.

The stateful `map` with a global variable is pretty hairy. I don't think we should recommend people do that.



Whenever I in the past wanted to calculate statistics for datasets I used `functools` similarly to how it's described in the colab, but I always felt it was a bit of a hassle to use it together with multiprocessing, which is why I picked up the issue, to do it "once and for all".

Should i close this and open another PR, with descriptions of how to use `map` for reduction, or?

Yes I think good documentation is the way to go here. @mariosasko 's examples are clear and efficient.

Maybe we could have an `Aggregations` section in the `Process` page with some guides on how to:
- use `.map()` to compute aggregates
- use `.with_format("arrow")` for max, min, etc. to save RAM and get max speed
- use a multiprocessed `.map()` to get partial results in parallel and combine them (max text length example)
- (advanced) use multiprocessing with an arbitrary accumulator (word count example)

And also a new conceptual guide on `Multiprocessed mapping` to say that it helps speed up CPU intensive processing but why it may lead to incorrect results when computing aggregates.

cc @stevhliu for visibility and if you have some comments

I would create a `Reduce` - to be more exact - subsection under `Map` to demonstrate these examples since we're showing how they can be done with the `Dataset.map` function. It'd also be good to add a link to the new concept guide from this section to solidify user understanding :)

Coolio. Ill close this PR and get going on another one adding what we've discussed during the next couple of days!

Is adding a section to the docs still planned? Couldn't find any related PR. 

There is a new integration with polars which is convenient btw. Here is an example for computing the length of the longest dialogue in a dataset using polars:

```python
>>> from datasets import load_dataset
>>> ds = load_dataset("HuggingFaceTB/smoltalk", "all", split="train")
>>> df = ds.to_polars()
>>> df.head()
shape: (5, 2)
┌─────────────────────────────────┬───────────────────┐
│ messages                        ┆ source            │
│ ---                             ┆ ---               │
│ list[struct[2]]                 ┆ str               │
╞═════════════════════════════════╪═══════════════════╡
│ [{"The function \( g(x) \) sat… ┆ numina-cot-100k   │
│ [{"Ben twice chooses a random … ┆ numina-cot-100k   │
│ [{"Find all values of $x$ that… ┆ numina-cot-100k   │
│ [{"How can you help me? I'm wr… ┆ smol-magpie-ultra │
│ [{"Extract and present the mai… ┆ smol-summarize    │
└─────────────────────────────────┴───────────────────┘
>>> df["messages"].list.len().max()
58
```

For very large scale dataset it can be worth using `map()` on batches of data to compute intermediate results, save some memory, and cache the result:

```python
>>> f = lambda df: pl.DataFrame({"messages_max_length": [df["messages"].list.len().max()]})
>>> intermediate_ds = ds.with_format("polars").map(f, batched=True)  # you can also set batch_size=
>>> intermediate_ds.to_polars()["messages_max_length"].max()
58
```
This last method can be used to implement a map + intermediate reduce + final reduce approach