class CustomFormatter(Formatter[dict, ColumnFormat, dict]):
    """
    A user-defined custom formatter function defined by a ``transform``.
    The transform must take as input a batch of data extracted for an arrow table using the python extractor,
    and return a batch.
    If the output batch is not a dict, then output_all_columns won't work.
    If the ouput batch has several fields, then querying a single column won't work since we don't know which field
    to return.
    """

    def __init__(self, transform: Callable[[dict], dict], features=None, token_per_repo_id=None, **kwargs):
        super().__init__(features=features, token_per_repo_id=token_per_repo_id)
        self.transform = transform

    def format_row(self, pa_table: pa.Table) -> dict:
        formatted_batch = self.format_batch(pa_table)
        try:
            return _unnest(formatted_batch)
        except Exception as exc:
            raise TypeError(
                f"Custom formatting function must return a dict of sequences to be able to pick a row, but got {formatted_batch}"
            ) from exc

    def format_column(self, pa_table: pa.Table) -> ColumnFormat:
        formatted_batch = self.format_batch(pa_table)
        if hasattr(formatted_batch, "keys"):
            if len(formatted_batch.keys()) > 1:
                raise TypeError(
                    "Tried to query a column but the custom formatting function returns too many columns. "
                    f"Only one column was expected but got columns {list(formatted_batch.keys())}."
                )
        else:
            raise TypeError(
                f"Custom formatting function must return a dict to be able to pick a row, but got {formatted_batch}"
            )
        try:
            return formatted_batch[pa_table.column_names[0]]
        except Exception as exc:
            raise TypeError(
                f"Custom formatting function must return a dict to be able to pick a row, but got {formatted_batch}"
            ) from exc

    def format_batch(self, pa_table: pa.Table) -> dict:
        batch = self.python_arrow_extractor().extract_batch(pa_table)
        batch = self.python_features_decoder.decode_batch(batch)
        return self.transform(batch)
/Users/nielsrogge/Documents/python_projecten/datasets/src/datasets/formatting/formatting.py