[Feature request] Add dataset.ragged_map() function for many-to-many transformations
Author: jarednielsen
State: closed

`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.



Comments: 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.

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!