[Feature request] Add dataset.ragged_map() function for many-to-many transformations
MEMBER
Actually `map(batched=True)` can already change the size of the dataset.
It can accept examples of length `N` and returns a batch of length `M` (can be null or greater than `N`).

I'll make that explicit in the doc that I'm currently writing.
2020-07-09T08:10:34Z
https://github.com/huggingface/datasets/issues/360#issuecomment-655977734
https://api.github.com/repos/huggingface/datasets/issues/360
MDEyOklzc3VlQ29tbWVudDY1NTk3NzczNA==
https://api.github.com/repos/huggingface/datasets/issues/comments/655977734/reactions
2020-07-09T08:10:34Z
https://api.github.com/repos/huggingface/datasets/issues/comments/655977734
https://avatars.githubusercontent.com/u/7353373?v=4
https://api.github.com/users/thomwolf/events{/privacy}
https://api.github.com/users/thomwolf/followers
https://api.github.com/users/thomwolf/following{/other_user}
https://api.github.com/users/thomwolf/gists{/gist_id}
https://github.com/thomwolf
thomwolf
MDQ6VXNlcjczNTMzNzM=
https://api.github.com/users/thomwolf/orgs
https://api.github.com/users/thomwolf/received_events
https://api.github.com/users/thomwolf/repos
https://api.github.com/users/thomwolf/starred{/owner}{/repo}
https://api.github.com/users/thomwolf/subscriptions
User
https://api.github.com/users/thomwolf
public
CONTRIBUTOR
You're two steps ahead of me :) In my testing, it also works if `M` < `N`.

A batched map of different length seems to work if you directly overwrite all of the original keys, but fails if any of the original keys are preserved.

For example,
```python
# Create a dummy dataset
dset = load_dataset("wikitext", "wikitext-2-raw-v1")["test"]
dset = dset.map(lambda ex: {"length": len(ex["text"]), "foo": 1})

# Do an allreduce on each batch, overwriting both keys
dset.map(lambda batch: {"length": [sum(batch["length"])], "foo": [1]})
# Dataset(schema: {'length': 'int64', 'foo': 'int64'}, num_rows: 5)

# Now attempt an allreduce without touching the `foo` key
dset.map(lambda batch: {"length": [sum(batch["length"])]})
# This fails with the error message below
```

```bash
  File "/path/to/nlp/src/nlp/arrow_dataset.py", line 728, in map
    arrow_schema = pa.Table.from_pydict(test_output).schema
  File "pyarrow/io.pxi", line 1532, in pyarrow.lib.Codec.detect
  File "pyarrow/table.pxi", line 1503, in pyarrow.lib.Table.from_arrays
  File "pyarrow/public-api.pxi", line 390, in pyarrow.lib.pyarrow_wrap_table
  File "pyarrow/error.pxi", line 85, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Column 1 named foo expected length 1 but got length 2
```

Adding the `remove_columns=["length", "foo"]` argument to `map()` solves the issue. Leaving the above error for future visitors. Perfect, thank you!
2020-07-09T19:28:10Z
https://github.com/huggingface/datasets/issues/360#issuecomment-656309230
https://api.github.com/repos/huggingface/datasets/issues/360
MDEyOklzc3VlQ29tbWVudDY1NjMwOTIzMA==
https://api.github.com/repos/huggingface/datasets/issues/comments/656309230/reactions
2020-07-09T19:31:29Z
https://api.github.com/repos/huggingface/datasets/issues/comments/656309230
https://avatars.githubusercontent.com/u/4564897?v=4
https://api.github.com/users/jarednielsen/events{/privacy}
https://api.github.com/users/jarednielsen/followers
https://api.github.com/users/jarednielsen/following{/other_user}
https://api.github.com/users/jarednielsen/gists{/gist_id}
https://github.com/jarednielsen
jarednielsen
MDQ6VXNlcjQ1NjQ4OTc=
https://api.github.com/users/jarednielsen/orgs
https://api.github.com/users/jarednielsen/received_events
https://api.github.com/users/jarednielsen/repos
https://api.github.com/users/jarednielsen/starred{/owner}{/repo}
https://api.github.com/users/jarednielsen/subscriptions
User
https://api.github.com/users/jarednielsen
public
`dataset.map()` enables one-to-one transformations. Input one example and output one example. This is helpful for tokenizing and cleaning individual lines.
`dataset.filter()` enables one-to-(one-or-none) transformations. Input one example and output either zero/one example. This is helpful for removing portions from the dataset.
However, some dataset transformations are many-to-many. Consider constructing BERT training examples from a dataset of sentences, where you map `["a", "b", "c"] -> ["a[SEP]b", "a[SEP]c", "b[SEP]c", "c[SEP]b", ...]`

I propose a more general `ragged_map()` method that takes in a batch of examples of length `N` and return a batch of examples `M`. This is different from the `map(batched=True)` method, which takes examples of length `N` and returns a batch of length `N`, processing individual examples in parallel. I don't have a clear vision of how this would be implemented efficiently and lazily, but would love to hear the community's feedback on this.

My specific use case is creating an end-to-end ELECTRA data pipeline. I would like to take the raw WikiText data and generate training examples from this using the `ragged_map()` method, then export to TFRecords and train quickly. This would be a reproducible pipeline with no bash scripts. Currently I'm relying on scripts like https://github.com/google-research/electra/blob/master/build_pretraining_dataset.py, which are less general.

