Extend QuestionAnsweringExtractive template to handle nested columns
State: closed
Labels: name: enhancement

Currently the `QuestionAnsweringExtractive` task template and `preprare_for_task` only support "flat" features. We should extend the functionality to cover QA datasets like:

* `iapp_wiki_qa_squad`
* `parsinlu_reading_comprehension`

where the nested features differ with those from `squad` and trigger an `ArrowNotImplementedError`:

```
---------------------------------------------------------------------------
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-12-50e5b8f69c20> in <module>
----> 1 ds.prepare_for_task("question-answering-extractive")[0]

~/git/datasets/src/datasets/arrow_dataset.py in prepare_for_task(self, task)
   1436         # We found a template so now flush `DatasetInfo` to skip the template update in `DatasetInfo.__post_init__`
   1437         dataset.info.task_templates = None
-> 1438         dataset = dataset.cast(features=template.features)
   1439         return dataset
   1440 

~/git/datasets/src/datasets/arrow_dataset.py in cast(self, features, batch_size, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, num_proc)
    977         format = self.format
    978         dataset = self.with_format("arrow")
--> 979         dataset = dataset.map(
    980             lambda t: t.cast(schema),
    981             batched=True,

~/git/datasets/src/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, desc)
   1600 
   1601         if num_proc is None or num_proc == 1:
-> 1602             return self._map_single(
   1603                 function=function,
   1604                 with_indices=with_indices,

~/git/datasets/src/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    176         }
    177         # apply actual function
--> 178         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    179         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    180         # re-apply format to the output

~/git/datasets/src/datasets/fingerprint.py in wrapper(*args, **kwargs)
    395             # Call actual function
    396 
--> 397             out = func(self, *args, **kwargs)
    398 
    399             # Update fingerprint of in-place transforms + update in-place history of transforms

~/git/datasets/src/datasets/arrow_dataset.py in _map_single(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, new_fingerprint, rank, offset, desc)
   1940                         )  # Something simpler?
   1941                         try:
-> 1942                             batch = apply_function_on_filtered_inputs(
   1943                                 batch,
   1944                                 indices,

~/git/datasets/src/datasets/arrow_dataset.py in apply_function_on_filtered_inputs(inputs, indices, check_same_num_examples, offset)
   1836                 effective_indices = [i + offset for i in indices] if isinstance(indices, list) else indices + offset
   1837             processed_inputs = (
-> 1838                 function(*fn_args, effective_indices, **fn_kwargs) if with_indices else function(*fn_args, **fn_kwargs)
   1839             )
   1840             if update_data is None:

~/git/datasets/src/datasets/arrow_dataset.py in <lambda>(t)
    978         dataset = self.with_format("arrow")
    979         dataset = dataset.map(
--> 980             lambda t: t.cast(schema),
    981             batched=True,
    982             batch_size=batch_size,

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.cast()

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.ChunkedArray.cast()

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/compute.py in cast(arr, target_type, safe)
    241     else:
    242         options = CastOptions.unsafe(target_type)
--> 243     return call_function("cast", [arr], options)
    244 
    245 

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/_compute.pyx in pyarrow._compute.call_function()

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/_compute.pyx in pyarrow._compute.Function.call()

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

~/miniconda3/envs/datasets/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: Unsupported cast from struct<answer_end: list<item: int32>, answer_start: list<item: int32>, text: list<item: string>> to struct using function cast_struct
```