Calling `cast_column`/`remove_columns` and a sequence of `map` operations ends up making `faiss` fail with `ValueError`
State: closed
Labels: description: Something isn't working
name: bug

First of all, sorry in advance for the unclear title, but this bug is weird to explain (at least for me), so I tried my best to summarize all the information in this issue.

## Describe the bug

Calling a certain combination of operations over a 🤗 `Dataset` and then trying to calculate the `faiss` index with `.add_faiss_index` ends up throwing an exception while trying to set the format back of a previously removed column. But this just happens over certain conditions... I'll present some scenarios below!

## Steps to reproduce the bug

Assuming the following dataset named `sample.csv` with some IMDb data:

```csv
id,title,summary
1877830,"The Batman","When a sadistic serial killer begins murdering key political figures in Gotham, Batman is forced to investigate the city's hidden corruption and question his family's involvement."
9419884,"Doctor Strange in the Multiverse of Madness","Doctor Strange teams up with a mysterious teenage girl from his dreams who can travel across multiverses, to battle multiple threats, including other-universe versions of himself, which threaten to wipe out millions across the multiverse. They seek help from Wanda the Scarlet Witch, Wong and others."
11138512,"The Northman","From visionary director Robert Eggers comes The Northman, an action-filled epic that follows a young Viking prince on his quest to avenge his father's murder." 
1745960,"Top Gun: Maverick","After more than thirty years of service as one of the Navy's top aviators, Pete Mitchell is where he belongs, pushing the envelope as a courageous test pilot and dodging the advancement in rank that would ground him."
```

We'll be able to reproduce the bug using the following piece of code:

```python
# Sample code to reproduce the bug
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
import torch

torch.set_grad_enabled(False)
ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

from datasets import load_dataset, Value

ds = load_dataset("csv", data_files=["sample.csv"], split="train")
ds = ds.cast_column("id", Value("int32")) # from `int64` to `int32`
ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])})
ds = ds.remove_columns(["title", "summary"])

def generate_embeddings(x):
    return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}

ds = ds.map(generate_embeddings)
ds = ds.remove_columns("inputs")
ds.add_faiss_index(column="embeddings") # It fails here!
```

The code above is an adaptation of https://huggingface.co/docs/datasets/faiss_es, for the sake of presenting the bug with a simple example.

## Expected results

Ideally, the `faiss` index should be calculated over the 🤗 `Dataset` and no exception should be triggered.

## Actual results

But what happens instead is that a `ValueError: Columns ['inputs'] not in the dataset. Current columns in the dataset: ['id', 'embeddings']`, which makes no sense as that column has been previously dropped.

## Environment info

<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.2.2
- Platform: Linux-5.4.0-1074-azure-x86_64-with-glibc2.31
- Python version: 3.9.5
- PyArrow version: 8.0.0
- Pandas version: 1.4.2


Comments:
It works if we either remove the `ds = ds.cast_column("id", Value("int32"))` line from the code above, or if instead calling `ds.remove_columns()` we remove the columns inside each mapping as `ds.map(..., remove_columns=[...])` instead of right after the mapping.

Both of those solutions seem to fix the issue, so the root cause of it may be around that. Sorry I cannot provide you more insights, in case I get to fix it I'll submit a PR, in the meanwhile the code that I'm using as a workaround is the following:

```python
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
import torch

torch.set_grad_enabled(False)
ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")

from datasets import load_dataset, Value

ds = load_dataset("csv", data_files=["sample.csv"], split="train")
ds = ds.cast_column("id", Value("int32"))
ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])}, remove_columns=["title", "summary"])

def generate_embeddings(x):
    return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}

ds = ds.map(generate_embeddings, remove_columns=["inputs"])
ds.add_faiss_index(column="embeddings")
```

FYI the main reason I want to use `dataset.remove_columns` rather than the function inside `dataset.map` is because according to the 🤗 Datasets documentation, it's faster.

"🤗 Datasets also has a [Dataset.remove_columns()](https://huggingface.co/docs/datasets/v2.2.1/en/package_reference/main_classes#datasets.Dataset.remove_columns) method that is functionally identical, but faster, because it doesn’t copy the data of the remaining columns."

More information at https://huggingface.co/docs/datasets/process#map

Here I'm presenting all the scenarios so that you can further investigate the issue:

- ✅ `cast_column` -> `map` with `remove_columns` -> `map` with `remove_columns` -> `add_faiss_index`

  ```python
  from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
  import torch
  
  torch.set_grad_enabled(False)
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  
  from datasets import load_dataset, Value
  
  ds = load_dataset("csv", data_files=["sample.csv"], split="train")
  ds = ds.cast_column("id", Value("int32"))
  ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])}, remove_columns=["title", "summary"])
  
  def generate_embeddings(x):
      return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}
  
  ds = ds.map(generate_embeddings, remove_columns=["inputs"])
  ds.add_faiss_index(column="embeddings")
  ```

- ❌ `cast_column` -> `map` -> `remove_columns` -> `map` -> `remove_columns` -> `add_faiss_index`

  ```python
  from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
  import torch
  
  torch.set_grad_enabled(False)
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  
  from datasets import load_dataset, Value
  
  ds = load_dataset("csv", data_files=["sample.csv"], split="train")
  ds = ds.cast_column("id", Value("int32"))
  ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])})
  ds = ds.remove_columns(["title", "summary"])
  
  def generate_embeddings(x):
      return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}
  
  ds = ds.map(generate_embeddings)
  ds = ds.remove_columns(["inputs"])
  ds.add_faiss_index(column="embeddings")
  ```

- ❌ `cast_column` -> `map` with `remove_columns` -> `map` -> `remove_columns` -> `add_faiss_index`


  ```python
  from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
  import torch
  
  torch.set_grad_enabled(False)
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  
  from datasets import load_dataset, Value
  
  ds = load_dataset("csv", data_files=["sample.csv"], split="train")
  ds = ds.cast_column("id", Value("int32"))
  ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])}, remove_columns=["title", "summary"])
  
  def generate_embeddings(x):
      return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}
  
  ds = ds.map(generate_embeddings)
  ds = ds.remove_columns(["inputs"])
  ds.add_faiss_index(column="embeddings")
  ```

- ✅ `cast_column` -> `map` -> `remove_columns` -> `map` with `remove_columns` -> `add_faiss_index`


  ```python
  from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
  import torch
  
  torch.set_grad_enabled(False)
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  
  from datasets import load_dataset, Value
  
  ds = load_dataset("csv", data_files=["sample.csv"], split="train")
  ds = ds.cast_column("id", Value("int32"))
  ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])})
  ds = ds.remove_columns(["title", "summary"])
  
  def generate_embeddings(x):
      return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}
  
  ds = ds.map(generate_embeddings, remove_columns=["inputs"])
  ds.add_faiss_index(column="embeddings")
  ```

- ✅ `map` -> `remove_columns` -> `map` -> `remove_columns` -> `add_faiss_index`


  ```python
  from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
  import torch
  
  torch.set_grad_enabled(False)
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
  
  from datasets import load_dataset, Value
  
  ds = load_dataset("csv", data_files=["sample.csv"], split="train")
  ds = ds.map(lambda x: {"inputs": f"{ctx_tokenizer.sep_token}".join(["title", "summary"])})
  ds = ds.remove_columns(["title", "summary"])
  
  def generate_embeddings(x):
      return {"embeddings": ctx_encoder(**ctx_tokenizer(x["inputs"], return_tensors="pt"))[0][0].numpy()}
  
  ds = ds.map(generate_embeddings)
  ds = ds.remove_columns(["inputs"])
  ds.add_faiss_index(column="embeddings")
  ```

So on, I've created #4411 so as to fix the bug with `remove_columns` under certain conditions before `add_faiss_index`, which means that the scenarios not working above are already working fine.