Issue: Add column type guessing from map return function
State: closed
Author: piercefreeman

As discussed [here](https://github.com/huggingface/datasets/issues/5965), there are some cases where datasets is unable to automatically promote columns during mapping. The fix is to explicitly provide a `features` definition so pyarrow can configure itself with the right column types from the outset.

This PR provides an alternative approach, which is functionally equivalent to specifying features but a bit cleaner within a larger mapping pipeline. It allows clients to typehint the return variable coming from the mapper function - if we find one of these type annotations specified, and no explicit features have been passed in, we'll try to convert it into a Features map. If the map function runs and casting is unable to succeed, it will raise a DatasetTransformationNotAllowedError that indicates the typehint may be to blame. It works for batched and non-batched mapping functions.

Currently supported column types:
- builtins primitives: string, int, float, bool
- dictionaries, lists (nested and one-deep)
- Optional types and None-Unions (synonymous with optional types)

It's used like:

```python
class DatasetTyped(TypedDict):
  texts: list[str]

def dataset_typed_map(batch) -> DatasetTyped:
  return {"texts": [text.split() for text in batch["raw_text"]]}

dataset = {"raw_text": ["", "This is a test", "This is another test"]}

with Dataset.from_dict(dataset) as dset:
  new_dataset = dset.map(
    dataset_typed_map,
    batched=True,
    batch_size=1,
    num_proc=1,
)
```

Open questions:
- Should logging indicate we have automatically guessed these types? Or proceed quietly until we hit an error (as is the current implementation).

Comments:
Thanks for working on this. However, having thought about this issue a bit more, supporting this doesn't seem like a good idea - it's better to be explicit than implicit, according to the Zen of Python 🙂. Also, I don't think many users would use this, so this raises the question of whether this is something we want to maintain.

cc @lhoestq for the 2nd opinion

@mariosasko I was going to quote the Zen of Python in the other direction :) To me, this actually is much more explicit than the current behavior of guessing pyarrow types based on the raw dictionary return values. Explicit typehinting is increasingly the de facto way to deal with this dynamic type serialization - plus it feels like a clearer fit to me than separating out the mapper function from the feature column definition in the call to the actual `.map()`. Another benefit is providing typehinting support for clients that use mypy or other static typecheckers to detect return mismatches.

But will leave it to you and @lhoestq to see if it's something you'd like in core versus a support package.

I meant that explicitly specifying the target features  (the `features` param) is cleaner (easier to track) than relying on type hints.

Passing features= to `map()` is richer and more explicit. Also I don't think users would guess that such API exist.

Other libraries like dask also infer the type from the output or requires the typing to be specified using the `meta` argument

Point about discoverability is a fair one, would certainly need some docs around it. All good! Will close this out and keep in our extension utilities.