[WIP] Adding Support for Reading Pandas Category
State: closed | Author: justin-yan (CONTRIBUTOR)

@lhoestq - continuing our conversation from https://github.com/huggingface/datasets/issues/1906#issuecomment-784247014

The goal of this PR is to support `Dataset.from_pandas(df)` where the dataframe contains a Category.

Just the 4 line change below actually does seem to work:

```
>>> from datasets import Dataset
>>> import pandas as pd
>>> df = pd.DataFrame(pd.Series(["a", "b", "c", "a"], dtype="category"))
>>> ds = Dataset.from_pandas(df)
>>> ds.to_pandas()
   0
0  a
1  b
2  c
3  a
>>> ds.to_pandas().dtypes
0    category
dtype: object
```

save_to_disk, etc. all seem to work as well.  The main things that are theoretically "incorrect" if we leave this are:

```
>>> ds.features.type
StructType(struct<0: int64>)
```
there are a decent number of references to this property in the library, but I can't find anything that seems to actually break as a result of this being int64 vs. dictionary?  I think the gist of my question is: a) do we *need* to change the dtype of Classlabel and have get_nested_type return a pyarrow.DictionaryType instead of int64? and b) do you *want* it to change?  The biggest challenge I see to implementing this correctly is that the data will need to be passed in along with the pyarrow schema when instantiating the Classlabel (I *think* this is unavoidable, since the type itself doesn't contain the actual label values) which could be a fairly intrusive change - e.g. `from_arrow_schema`'s interface would need to change to include optional arrow data?  Once we start going down this path of modifying the public interfaces I am admittedly feeling a little bit outside of my comfort zone

Additionally I think `int2str`, `str2int`, and `encode_example` probably won't work - but I can't find any usages of them in the library itself.

Comments: Thanks ! could you maybe add a few tests in test_arrow_dataset.py to make sure from_pandas works as expected with categorical types ?

In particular I'm pretty sure that if you now try to `cast` the dataset to the same features at its current features, it will break instead of just being a no-op.
This is because `features.type` returns an arrow int64 type for the classlabel column instead of the arrow dictionary type that you have in the arrow table. There are two issues in this case:
- it will try to replace the arrow type from dictionary to int64 instead of being a no-op
- it will crash because pyarrow is not able to cast a dictionary to int64 (even if it's actually possible do cast the column by hand by accessing the sub-array of the dictionary array containing the indices/integers)

It would be awesome to fix this case ! Ideally the arrow `pa_type` of classlabel ([here](https://github.com/huggingface/datasets/blob/7072e1becd69d421d863374b825e3da4c6551798/src/datasets/features.py#L558)) should be an arrow dictionary type. This should fix the issue. Then we can start working on backward compatibility.

Let me know if you have questions or if I can help.
In particular if there is some glue-ing to do I can take care of that if you want ;)

--------------

Also just a few information regarding the functions you mentioned

`int2str` and  `str2int` are used by users to transforms the labels if they want to. Here sine ClassLabel is instantiated without the class names, they would crash. I was about to make a PR to disallow the creation of an empty ClassLabel feature type.
Therefore can you provide class_names= when creating the ClassLabel ?

`encode_example` is mostly used with a dataset builder (e.g. squad.py) so it's not used when using .from_pandas.




Got it - that's super helpful, I was trying to figure out what would break!

I think there are two issues we're discussing here:

1. modifying the pa_type of ClassLabel: totally agree with you on that one if that's OK from a back-compat perspective.  (i.e. are users of `datasets` not supposed to access or use the .pa_type attribute of ClassLabel?)
2. creating a ClassLabel requires information that's not present on the pa.DictionaryType object:  I think the crux of the problem is that at this line (https://github.com/huggingface/datasets/pull/1936/files#diff-54081ede051fd0a7ef65748c481cc06f90209f01bb89968747089d13a2ca052bR933) - you only have access to the `pa_type`, which is `DictionaryType[int8, string]`. I've unpacked it and looked at all of the available methods, and I don't believe that any of the actual values ("names") are present - those are stored on the `pyarrow.DictArray.dictionary` attribute (i.e. as data, not on the pyarrow.DataType) - so in order to actually be able to instantiate the ClassLabel with the names= parameter, we need to pass in more information to this method.

We *could* mostly accomplish this by modifying https://github.com/huggingface/datasets/pull/1936/files#diff-54081ede051fd0a7ef65748c481cc06f90209f01bb89968747089d13a2ca052bR909 to accept a pyarrow Table in addition to the type, and it's not too difficult to do, but it feels a little bit off to me:

- It feels a bit off that a "schema" definition will change depending on what data gets added to the dataset.  In particular, if someone adds rows or concatenates two datasets, the ClassLabel "names" will also need to change, right?  I think maybe we're getting around this because a Dataset is immutable (I think?) and so any new dataset is freshly constructed, but for example - I think this check wouldn't work for `ClassLabel`s if we were to compare the `Dataset.features` instead of the underlying pyarrow type https://github.com/huggingface/datasets/blob/master/src/datasets/arrow_dataset.py#L2664
- To that end I wonder if ClassLabel should actually just be the "type" akin to Category, and the "names" should be considered "data" and not part of the "type"? Similar to how pyarrow maintains two data objects - the array of indices and the array of string values.

With that in mind, I'm wondering if you *should* allow an empty ClassLabel (and`int2str`, etc. can be updated to have more descriptive error messages if labels aren't provided or inferred), and if the underlying data is a pa.DictionaryType, then the names can be inferred and applied at these points in the code:
- https://github.com/huggingface/datasets/blob/96578adface7e4bc1f3e8bafbac920d72ca1ca60/src/datasets/arrow_dataset.py#L274
- https://github.com/huggingface/datasets/blob/96578adface7e4bc1f3e8bafbac920d72ca1ca60/src/datasets/arrow_dataset.py#L686
- https://github.com/huggingface/datasets/blob/96578adface7e4bc1f3e8bafbac920d72ca1ca60/src/datasets/arrow_dataset.py#L673

I think perhaps the mismatch here is when the data is stored on disk as an int there should be a convenient way of saying "this is a dictionary and here are some explicitly provided labels", whereas when it's stored as a string, we'd ideally like to say "this is a Category and please condense the representation and automatically infer the labels".

Sorry for the long comment!  Hopefully my thoughts make sense - thanks for taking the time to discuss!

Yes that makes sense. I completely forgot that the label names of an arrow Dictionary type were not stored in the type but in the DictionaryArray.

This is made me realize that it's actually pretty unpractical and I feel that handling this can add unnecessary complexity in the handling of dtypes.
More specifically:
- it's not possible to create a DictionaryArray from a call to pyarrow.array with python objects, which is the function we use to convert python objects to pyarrow objects (or we would need to convert the python objects to pandas categorical series beforehand but it doesn't work for nested types)
- casting nested types containing Dictionary types would require a lot of array manipulations since it's not compatible with pyarrow.array.cast

I feel like the original feature request (support of pandas Categorical) should be addressable without adding so much complexity to the library.

If we admit that we don't want to deal with arrow Dictionary type, maybe we can simply convert the pandas categorical series to an int64 series and set the feature type to the right ClassLabel in `from_pandas`. We can have the reverse operation in `to_pandas`. This way we don't need to support the arrow DictionaryType and so we can keep simple/accessible code for conversion from python to arrow and also for type casting. Let me know what you think.

In the future depending on the usage of the ClassLabel types with pandas/pyarrow we might reconsider this but for now I believe this simple solution is enough.

I like that idea!  Let me try working up a PR for this

OK! I just whipped up the `from_pandas()` portion of this PR, and it works, though I'm not *super* familiar with the available APIs so I'm not sure if there's a more "vectorized" way of doing all of these updates - so happy to get some feedback and iterate!

Apologies for multiple commits - I realized how to solve a few different problems right after I gave up and pushed with the intent to ask for help :-)

I wanted to get some guidance on how to handle the reverse direction: I think there are two main areas to look at, `.to_pandas()` and also `.set_format('pandas')` and then pulling out a dataframe like so: `dataset[:]`.  Is there a single place where I can handle both of these cases at once or do these need to be handled independently?

Thanks ! This is awesome :) 
Could you also add a test ? There is already `test_to_pandas` in test_arrow_dataset.py
Feel free to complete this test to make sure it works for Categorical :)

To make it work with the "pandas" formating (when you do `set_format("pandas")` and then query `dataset[0]`, `dataset[:]`, etc.), you can take a look and the `PandasFormatter` in formatting.py
It takes a pyarrow table as input of its formatting methods (one method for rows, one for columns and one for batches) and returns a pandas DataFrame (or a Series for the method for formatting a column). You can cast to Categorical in each one of the formatter methods and it should work directly when you use a pandas-formatted dataset.

This formatter can then also be used in `to_pandas` (currently it does `pa_table.to_pandas()` but `PandasFormatter().format_batch(pa_table)` can be used instead).