Title: KeyError on using map after renaming a column
Author: gchhablani (CONTRIBUTOR)
State: closed
Created: 2021-03-10T18:54:17Z

Hi,

I'm trying to use `cifar10` dataset. I want to rename the `img` feature to `image` in order to make it consistent with `mnist`, which I'm also planning to use. By doing this, I was trying to avoid modifying `prepare_train_features` function.

Here is what I try:

```python
transform = Compose([ToPILImage(),ToTensor(),Normalize([0.0,0.0,0.0],[1.0,1.0,1.0])])
def prepare_features(examples):
    images = []
    labels = []
    print(examples)
    for example_idx, example in enumerate(examples["image"]):
        if transform is not None:
            images.append(transform(examples["image"][example_idx].permute(2,0,1)))
        else:
            images.append(examples["image"][example_idx].permute(2,0,1))
        labels.append(examples["label"][example_idx])
    output = {"label":labels, "image":images}
    return output

raw_dataset = load_dataset('cifar10')
raw_dataset.set_format('torch',columns=['img','label'])
raw_dataset = raw_dataset.rename_column('img','image')

features = datasets.Features({
            "image": datasets.Array3D(shape=(3,32,32),dtype="float32"),
            "label": datasets.features.ClassLabel(names=[
                            "airplane",
                            "automobile",
                            "bird",
                            "cat",
                            "deer",
                            "dog",
                            "frog",
                            "horse",
                            "ship",
                            "truck",
                        ]),
        })
train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)
```
The error:
```python
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-54-bf29672c53ee> in <module>()
     14                         ]),
     15         })
---> 16 train_dataset = raw_dataset.map(prepare_features, features = features,batched=True, batch_size=10000)

2 frames
/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1287         test_inputs = self[:2] if batched else self[0]
   1288         test_indices = [0, 1] if batched else 0
-> 1289         update_data = does_function_return_dict(test_inputs, test_indices)
   1290         logger.info("Testing finished, running the mapping function on the dataset")
   1291 

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in does_function_return_dict(inputs, indices)
   1258             fn_args = [inputs] if input_columns is None else [inputs[col] for col in input_columns]
   1259             processed_inputs = (
-> 1260                 function(*fn_args, indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
   1261             )
   1262             does_return_dict = isinstance(processed_inputs, Mapping)

<ipython-input-52-b4dccbafb70d> in prepare_features(examples)
      3     labels = []
      4     print(examples)
----> 5     for example_idx, example in enumerate(examples["image"]):
      6         if transform is not None:
      7             images.append(transform(examples["image"][example_idx].permute(2,0,1)))

KeyError: 'image'
```

The print statement inside returns this:
```python
{'label': tensor([6, 9])}
```
Apparently, both `img` and `image` do not exist after renaming. 

Note that this code works fine with `img` everywhere.

Notebook: https://colab.research.google.com/drive/1SzESAlz3BnVYrgQeJ838vbMp1OsukiA2?usp=sharing

Comments:
Hi,

Actually, the error occurs due to these two lines:
```python
raw_dataset.set_format('torch',columns=['img','label'])
raw_dataset = raw_dataset.rename_column('img','image')
```
`Dataset.rename_column` doesn't update the `_format_columns` attribute, previously defined by `Dataset.set_format`, with a new column name which is why this new column is missing in the output.

Hi @mariosasko,

Thanks for opening a PR on this :)
Why does the old name also disappear?

I just merged a @mariosasko 's PR that fixes this issue.
If it happens again, feel free to re-open :)