Title: Batched map fails when removing all columns
Author: villmow (NONE)
State: closed
Created: 2021-04-16T11:17:01Z
Labels: bug

Hi @lhoestq ,

I'm hijacking this issue, because I'm currently trying to do the approach you recommend:

> Currently the optimal setup for single-column computations is probably to do something like
> 
> ```python
> result = dataset.map(f, input_columns="my_col", remove_columns=dataset.column_names)
> ```

Here is my code: (see edit, in which I added a simplified version

```
This is the error:
```bash
pyarrow.lib.ArrowInvalid: Column 1 named tokens expected length 8964 but got length 1000
```
I wonder why this error occurs, when I delete every column? Can you give me a hint?

### Edit:
I preprocessed my dataset before (using map with the features argument) and saved it to disk. May this be part of the error?  I can iterate over the
complete dataset and print every sample before calling map. There seems to be no other problem with the dataset.

I tried to simplify the code that crashes:

```python
# works
log.debug(dataset.column_names)
log.debug(dataset)
for i, sample in enumerate(dataset):
    log.debug(i, sample)

# crashes
counted_dataset = dataset.map(
    lambda x: {"a": list(range(20))},
    input_columns=column,
    remove_columns=dataset.column_names,
    load_from_cache_file=False,
    num_proc=num_workers,
    batched=True,
)
```

```
pyarrow.lib.ArrowInvalid: Column 1 named tokens expected length 20 but got length 1000
```

Edit2: 

May this be a problem with a schema I set when preprocessing the dataset before? I tried to add the `features` argument to the function and then I get a new error:

```python
# crashes
counted_dataset = dataset.map(
    lambda x: {"a": list(range(20))},
    input_columns=column,
    remove_columns=dataset.column_names,
    load_from_cache_file=False,
    num_proc=num_workers,
    batched=True,
    features=datasets.Features(
        {
              "a": datasets.Sequence(datasets.Value("int32"))
         }
    )
)
```

```
 File "env/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1704, in _map_single
    writer.write_batch(batch)
  File "env/lib/python3.8/site-packages/datasets/arrow_writer.py", line 312, in write_batch
    col_type = schema.field(col).type if schema is not None else None
  File "pyarrow/types.pxi", line 1341, in pyarrow.lib.Schema.field
KeyError: 'Column tokens does not exist in schema'
```

_Originally posted by @villmow in https://github.com/huggingface/datasets/issues/2193#issuecomment-820230874_

Comments:
I found the problem. I called `set_format` on some columns before. This makes it crash. Here is a complete example to reproduce:

```python
from datasets import load_dataset
sst = load_dataset("sst")
sst.set_format("torch", columns=["label"], output_all_columns=True)
ds = sst["train"]

# crashes
ds.map(
    lambda x: {"a": list(range(20))},
    remove_columns=ds.column_names,
    load_from_cache_file=False,
    num_proc=1,
    batched=True,
)
```

Thanks for reporting and for providing this code to reproduce the issue, this is really helpful !

I merged a fix, it should work on `master` now :)
We'll do a new release soon !