Hey Sam, Did you consider DictionaryArray? (https://arrow.apache.org/docs/python/data.html#dictionary-arrays) It's to_pandas will return pd.Categorical.
Rok On Wed, Jan 5, 2022 at 3:35 PM Sam Davis <[email protected]> wrote: > > Hi, > > I'm looking at defining a schema for a table where one of the values is > inherently categorical/enumerable and we're ultimately ending up loading it > as a Pandas DataFrame. I cannot seem to find a decent way of achieving this. > > For example, the column may always be known to contain the values ["a", "b", > "c", "d"]. Stating this as a stringly-typed column in the schema is a bad > idea as it permits all strings and requires more storage than necessary for > longer strings, stating it as an integer column is a bad idea as you lose > context and force the user to cast after loading, and the dictionary type > does not allow you to specify the values in the schema so similarly loses all > meaning. > > I have been playing with the API all morning and from what I can tell there > is no easy way of achieving this. Am I missing something obvious? > > --- > > One possible route I thought of is to define an extension type and then > implement the `to_pandas_dtype` method. Yes this method permits all known > values whilst in Arrow-land, but it at least documents the known type and, so > I thought, any values not within the `to_pandas_dtype` return will be set to > null on conversion anyway. > > However, this seems to require unnecessarily special-casing a whole bunch of > code to handle extension types. e.g. just creating a scalar of this type > requires using a different API. It seems like `pa.scalar` should be able to > work this out? This example defines a wrapper for int32, and then tries to > create a scalar of this type showing that the user has to call a special > method rather than just the normal API: > > ``` > import pyarrow as pa > > > class IntegerWrapper(pa.ExtensionType): > > def __init__(self): > pa.ExtensionType.__init__(self, pa.int32(), "integer_wrapper") > > def __arrow_ext_serialize__(self): > # since we don't have a parameterized type, we don't need extra > # metadata to be deserialized > return b'' > > @classmethod > def __arrow_ext_deserialize__(self, storage_type, serialized): > # return an instance of this subclass given the serialized > # metadata. > return IntegerWrapper() > > > iw_type = IntegerWrapper() > > pa.register_extension_type(iw_type) > > # throws `ArrowNotImplementedError` > # pa.scalar(0, iw_type) > > # user must do this, but code should be able to do this? > pa.ExtensionScalar.from_storage(iw_type, pa.scalar(0, iw_type.storage_type)) > ``` > > and I can't seem to get the `to_pandas_dtype` to actually work for a wrapped > dictionary. e.g. > > ``` > import pyarrow as pa > > > class DictWrapper(pa.ExtensionType): > > def __init__(self): > pa.ExtensionType.__init__(self, pa.dictionary(pa.int8(), > pa.string()), "dict_wrapper") > > def __arrow_ext_serialize__(self): > # since we don't have a parameterized type, we don't need extra > # metadata to be deserialized > return b'' > > @classmethod > def __arrow_ext_deserialize__(self, storage_type, serialized): > # return an instance of this subclass given the serialized > # metadata. > return DictWrapper() > > def to_pandas_dtype(self): > from pandas.api.types import CategoricalDtype > return CategoricalDtype(categories=["a", "b"]) > > dw_type = DictWrapper() > > pa.register_extension_type(dw_type) > > arr = pa.ExtensionArray.from_storage( > dw_type, > pa.array(["a", "b", "c"], dw_type.storage_type) > ) > > arr > > arr.to_pandas() > > arr.to_pandas(categories=dw_type.to_pandas_dtype().categories.values) > ``` > > Best, > > Sam > IMPORTANT NOTICE: The information transmitted is intended only for the person > or entity to which it is addressed and may contain confidential and/or > privileged material. Any review, re-transmission, dissemination or other use > of, or taking of any action in reliance upon, this information by persons or > entities other than the intended recipient is prohibited. If you received > this in error, please contact the sender and delete the material from any > computer. Although we routinely screen for viruses, addressees should check > this e-mail and any attachment for viruses. We make no warranty as to absence > of viruses in this e-mail or any attachments.
