ArrowInvalid issue for squad v2 dataset

Hello, I am using the huggingface official question answering example notebook (https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/question_answering.ipynb). 

In the prepare_validation_features function, I made some modifications to tokenize a new set of quesions with the original contexts and save them in three different list called candidate_input_dis, candidate_attetion_mask and candidate_token_type_ids. When I try to run the next cell for dataset.map, I got the following error:

`ArrowInvalid: Column 1 named candidate_attention_mask expected length 1180 but got length 1178`

My code is as follows:

```
def generate_candidate_questions(examples):
  val_questions = examples["question"]
  candididate_questions = random.sample(datasets["train"]["question"], len(val_questions))
  candididate_questions = [x[:max_length] for x in candididate_questions]
  return candididate_questions

def prepare_validation_features(examples, use_mixing=False):
  pad_on_right = tokenizer.padding_side == "right"
  tokenized_examples = tokenizer(
      examples["question" if pad_on_right else "context"],
      examples["context" if pad_on_right else "question"],
      truncation="only_second" if pad_on_right else "only_first",
      max_length=max_length,
      stride=doc_stride,
      return_overflowing_tokens=True,
      return_offsets_mapping=True,
      padding="max_length",
  )
  if use_mixing:
    candidate_questions = generate_candidate_questions(examples)
    tokenized_candidates = tokenizer(
        candidate_questions if pad_on_right else examples["context"],
        examples["context"] if pad_on_right else candidate_questions,
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_length,
        stride=doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

  sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")

  tokenized_examples["example_id"] = []

  if use_mixing:
    tokenized_examples["candidate_input_ids"] = tokenized_candidates["input_ids"]
    tokenized_examples["candidate_attention_mask"] = tokenized_candidates["attention_mask"]
    tokenized_examples["candidate_token_type_ids"] = tokenized_candidates["token_type_ids"]

  for i in range(len(tokenized_examples["input_ids"])):
      sequence_ids = tokenized_examples.sequence_ids(i)
      context_index = 1 if pad_on_right else 0

      sample_index = sample_mapping[i]
      tokenized_examples["example_id"].append(examples["id"][sample_index])

      tokenized_examples["offset_mapping"][i] = [
          (o if sequence_ids[k] == context_index else None)
          for k, o in enumerate(tokenized_examples["offset_mapping"][i])
      ]

  return tokenized_examples



validation_features = datasets["validation"].map(
    lambda xs: prepare_validation_features(xs, True),
    batched=True,
    remove_columns=datasets["validation"].column_names
)
```

I guess this might happen because of the batched=True. I see similar issues in this repo related to arrow table length mismatch error, but in their cases, the numbers vary a lot. In my case, this error always happens when the expected length and unexpected length are very close. Thanks for the help!

Comments:
Hi ! This error happens when you use `map` in batched mode and then your function doesn't return the same number of values per column.

Indeed since you're using `map` in batched mode, `prepare_validation_features` must take a batch as input (i.e. a dictionary of multiple rows of the dataset), and return a batch.

However it seems like `tokenized_examples` doesn't have the same number of elements in each field. One field seems to have `1180` elements while `candidate_attention_mask` only has `1178`.