ENH: allow JSON (de)serialization of ExtensionDtypes
State: closed


 - xref #20612
 - closes #44705
- [x] tests added / passed
- [x] ensure all linting tests pass, see [here](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit) for how to run them

This PR allows ExtensionDtypes to be (de)serialized using `.to_json(orient='table')`.

Example:
```py
import datetime as dt
from typing import Any, Callable, Optional, Sequence
from typing import Tuple
from typing import Union
from typing import cast

import numpy as np
from pandas.api.extensions import ExtensionArray
from pandas.api.extensions import ExtensionDtype
from pandas.api.types import pandas_dtype
from pandas.core.dtypes.dtypes import register_extension_dtype
from pandas._typing import Dtype


@register_extension_dtype
class DateDtype(ExtensionDtype):
    @property
    def type(self):
        return dt.date

    @property
    def name(self):
        return "DateDtype"

    @classmethod
    def construct_from_string(cls, string: str):
        if not isinstance(string, str):
            raise TypeError(
                f"'construct_from_string' expects a string, got {type(string)}"
            )

        if string == cls.__name__:
            return cls()
        else:
            raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")

    @classmethod
    def construct_array_type(cls):
        return DateArray

    @property
    def na_value(self):
        return dt.date.min

    def __repr__(self) -> str:
        return self.name


class DateArray(ExtensionArray):
    def __init__(
        self,
        dates: Union[
            Sequence[dt.date], Tuple[np.ndarray, np.ndarray, np.ndarray], np.ndarray
        ],
    ) -> None:
        ldates = len(dates)
        if isinstance(dates, list):
            # pre-allocate the arrays since we know the size before hand
            self._year = np.zeros(ldates, dtype=np.uint16)  # 65535 (0, 9999)
            self._month = np.zeros(ldates, dtype=np.uint8)  # 255 (1, 31)
            self._day = np.zeros(ldates, dtype=np.uint8)  # 255 (1, 12)
            # populate them
            for i, (y, m, d) in enumerate(
                map(lambda date: (date.year, date.month, date.day), dates)
            ):
                self._year[i] = y
                self._month[i] = m
                self._day[i] = d

        elif isinstance(dates, tuple):
            # only support triples
            if ldates != 3:
                raise ValueError("only triples are valid")
            # check if all elements have the same type
            if not isinstance(dates[0], np.ndarray):
                raise TypeError(
                    f"tuple 0-element has an invalid type: {type(dates[0])}"
                )
            if not isinstance(dates[1], np.ndarray):
                raise TypeError(
                    f"tuple 1-element has an invalid type: {type(dates[1])}"
                )
            if not isinstance(dates[2], np.ndarray):
                raise TypeError(
                    f"tuple 2-element has an invalid type: {type(dates[2])}"
                )

            ly, lm, ld = (len(cast(np.ndarray, d)) for d in dates)
            if not ly == lm == ld:
                raise ValueError(
                    f"tuple members must have the same length: {(ly, lm, ld)}"
                )

            self._year = dates[0].astype(np.uint16)
            self._month = dates[1].astype(np.uint8)
            self._day = dates[2].astype(np.uint8)

        elif isinstance(dates, np.ndarray) and dates.dtype == "U10":
            self._year = np.zeros(ldates, dtype=np.uint16)  # 65535 (0, 9999)
            self._month = np.zeros(ldates, dtype=np.uint8)  # 255 (1, 31)
            self._day = np.zeros(ldates, dtype=np.uint8)  # 255 (1, 12)

            for i, dstr in np.ndenumerate(np.char.split(dates, sep="-")):
                self._year[i] = int(dstr[0])
                self._month[i] = int(dstr[1])
                self._day[i] = int(dstr[2])

        else:
            raise TypeError(f"{type(dates)} is not supported")

    @property
    def dtype(self) -> ExtensionDtype:
        return DateDtype()

    def astype(self, dtype, copy=True):
        dtype = pandas_dtype(dtype)

        if isinstance(dtype, DateDtype):
            data = self.copy() if copy else self
        else:
            # NOTE(duarte,30/11/21): na_value is probably wrong
            data = self.to_numpy(dtype=dtype, copy=copy, na_value=dt.date.min)

        return data

    @property
    def nbytes(self) -> int:
        return self._year.nbytes + self._month.nbytes + self._day.nbytes

    def __len__(self) -> int:
        return len(self._year)  # all 3 arrays are enforced to have the same length

    def __getitem__(self, item) -> Union[dt.date, "DateArray"]:
        if isinstance(item, int):
            return dt.date(self._year[item], self._month[item], self._day[item])
        else:
            raise NotImplementedError("only ints are supported as indexes")

    def __setitem__(self, key: Union[int, np.ndarray], value) -> None:
        if not isinstance(key, int):
            raise NotImplementedError("only ints are supported as indexes")

        if not isinstance(value, dt.date):
            raise TypeError("you can only set datetime.date types")

        self._year[key] = value.year
        self._month[key] = value.month
        self._day[key] = value.day

    def __repr__(self) -> str:
        return f"DateArray{list(zip(self._year, self._month, self._day))}"

    def copy(self) -> "DateArray":
        return DateArray((self._year.copy(), self._month.copy(), self._day.copy()))

    def isna(self) -> np.ndarray:
        return np.logical_and(
            np.logical_and(
                self._year == dt.date.min.year, self._month == dt.date.min.month
            ),
            self._day == dt.date.min.day,
        )

    @classmethod
    def _from_sequence(cls, scalars, *, dtype: Optional[Dtype] = None, copy=False):
        if isinstance(scalars, dt.date):
            pass
        elif isinstance(scalars, DateArray):
            pass
        elif isinstance(scalars, np.ndarray):
            scalars = scalars.astype("U10")  # 10 chars for yyyy-mm-dd
            return DateArray(scalars)
```

```py
import pandas as pd
import datetime as dt

da = DateArray([dt.date.today()])
df = pd.DataFrame({'da': da})
df_rtt = pd.read_json(df.to_json(orient='table'), orient='table')
# df.equals(df_rtt) won't work because:
# * the indexes are different (already happened)
# * the DateArray class only implements ints as indexes (its only a PoC)
assert df.dtypes == df_rtt.dtypes 
```


Labels:
name: Enhancement
name: IO JSON
description: read_json, to_json, json_normalize
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.

Comments:
@jmg-duarte Can you do the following:

1. Add a test for this new functionality.  You might want to follow the pattern in `pandas/test/io/json/test_json_table_schema.py`, but create a new file that uses your example extension type (date/time) and also do something with `Decimal` as well.  You can import an extension dtype implementation for `Decimal` in the package `pandas.tests.extension.decimal`.  
2. Update the documentation for `read_table()` to discuss this new functionality.
3. Update the what's new document to reflect this new functionality
> 2. Update the documentation for `read_table()` to discuss this new functionality.

Do you mean `read_json()`?

> > 2. Update the documentation for `read_table()` to discuss this new functionality.
> 
> Do you mean `read_json()`?

Yes.  Also, `to_json()` and the docs corresponding to https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html?highlight=read_json#table-schema, and more generally https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html?highlight=read_json#json

We'll need to state that for user-implemented extension arrays, we're possibly generating JSON Table schema that is non-conforming.  (e.g., if you do the test for `Decimal`, I would anticipate that would happen).  It's probably worth incrementing the schema version as well.

One other note - I see that tests are already failing.   You might want to run the existing tests for the table schema locally before pushing your next commit:
```
pytest -n 4 pandas/tests/io/json/test_json_table_schema.py
```
(The `-n 4` argument runs the tests in parallel to give you faster results)

A little checklist to track the current progress:
- [x] Add tests to `pandas/io/json/test_json_table_schema_ext_dtype.py`
  - [x] Include `Decimal` in the tests
  - [x] Pass tests
- [x] Update the docs
  - [x] `to_json()`
  - [x] `read_json()`
- [x] Increment the schema version
  - [x] Move `TABLE_SCHEMA_VERSION` to `_table_schema.py`
@attack68 You might be the best person to review this PR once the checklist at https://github.com/pandas-dev/pandas/pull/44722#issuecomment-984818057 is complete.  If not, do you have other suggestions?
I just tested `Decimal` on master and while it works, it does not make the roundtrip.
```py
In [1]: import pandas as pd

In [2]: from pandas.tests.extension.decimal import DecimalArray

In [3]: from decimal import Decimal

In [4]: pd.read_json(pd.DataFrame({'dec': DecimalArray([Decimal(10)])}).to_json(orient='table'), orient='table')
Out[4]: 
    dec
0  10.0

In [5]: pd.read_json(pd.DataFrame({'dec': DecimalArray([Decimal(10)])}).to_json(orient='table'), orient='table').dtypes
Out[5]: 
dec    object
dtype: object

In [6]: pd.DataFrame({'dec': DecimalArray([Decimal(10)])}).to_json(orient='table')
Out[6]: '{"schema":{"fields":[{"name":"index","type":"integer"},{"name":"dec","type":"string"}],"primaryKey":["index"],"pandas_version":"0.20.0"},"data":[{"index":0,"dec":10.0}]}'
```

Out of the box, the solution proposed in this PR will throw the following error:
```py
TypeError: All values must be of type <class 'decimal.Decimal'>
```
Due to the following code:
https://github.com/pandas-dev/pandas/blob/e99296998fc1118888164b71968c5f091da206e2/pandas/tests/extension/decimal/array.py#L89-L91

With a bit of adapting I can make this work, however I'm open to feedback on how to handle this.
I've updated the `to_json()` documentation, I don't think `read_json()` needs it but I don't mind adding a reference to which method the user needs to override to tap into the serialization. Let me know your thoughts.

On another note, why don't these lines use the `TABLE_SCHEMA_VERSION` variable?
https://github.com/pandas-dev/pandas/blob/e99296998fc1118888164b71968c5f091da206e2/pandas/io/json/_table_schema.py#L293-L294
> With a bit of adapting I can make this work, however I'm open to feedback on how to handle this.

It's why I suggested using  `ExtensionArray._from_sequence_of_strings()` for deserialization, `ExtensionArray._formatter()` for serialization, because you would then just put the string representation in the JSON, and those should be inverses of each other, and independent of the extension array implementations, you know the string representations of the underlying objects are there.

> On another note, why don't these lines use the `TABLE_SCHEMA_VERSION` variable?

Seems like this was never noticed before.  Can you move `TABLE_SCHEMA_VERSION` out of `_json.py` and into `_table_schema.py` as part of this PR?

Note in the docs in `_json.py` and `_table_schema.py`, there are example outputs shown that have that number in it, so when you change the version, you should also change the examples as well.
> I've updated the `to_json()` documentation, I don't think `read_json()` needs it but I don't mind adding a reference to which method the user needs to override to tap into the serialization. Let me know your thoughts.

I think in the overview pages I referred to earlier is where you want to discuss how `ExtensionArray` is handled with respect to the table schema.  For the API docs, you just need to indicate that if the `DataFrame` has a `Series` backed by a user-created `ExtensionArray`, then the output will not conform to the table schema standard because pandas will be outputting types that are not part of that standard, but that round-tripping (i.e., `read_json(df.to_json(orient="table"), orient="table")` will work.
> It's why I suggested using `ExtensionArray._from_sequence_of_strings()` for deserialization, `ExtensionArray._formatter()` for serialization, because you would then just put the string representation in the JSON, and those should be inverses of each other, and independent of the extension array implementations, you know the string representations of the underlying objects are there.

Due to my lack of experience with pandas (and a bit of documentation too) I'm doing this in an ad-hoc approach - *implement a bit, see what breaks and repeat*. 
I'll put breakpoints into the methods you suggest and work from there.

> Seems like this was never noticed before. Can you move `TABLE_SCHEMA_VERSION` out of `_json.py` and into `_table_schema.py` as part of this PR?
> 
> Note in the docs in `_json.py` and `_table_schema.py`, there are example outputs shown that have that number in it, so when you change the version, you should also change the examples as well.

Sure thing, I will address it.

> I think in the overview pages I referred to earlier is where you want to discuss how `ExtensionArray` is handled with respect to the table schema. For the API docs, you just need to indicate that if the `DataFrame` has a `Series` backed by a user-created `ExtensionArray`, then the output will not conform to the table schema standard because pandas will be outputting types that are not part of that standard, but that round-tripping (i.e., `read_json(df.to_json(orient="table"), orient="table")` will work.

I hadn't noticed the links were for "*extra*" documentation, I will also address this.
@jmg-duarte One other note, since this is your first PR.  I don't think you should keep pushing each commit you do, as it fires off our continuous integration testing scripts, and you know that you are still working on things.  Make sure the testing scripts work locally on your own machine, and when that's true, then you can push, and the full suite of tests will be run, possibly catching errors in different environments that don't happen on your own machine. When you make the doc changes, build them locally to make sure they look right, before pushing.
@Dr-Irv all tests are passing besides the docs, could you (or someone else) help me figure out the problem? The changes I made built locally
> @Dr-Irv all tests are passing besides the docs, could you (or someone else) help me figure out the problem? The changes I made built locally

I created a new issue for the CI failure:  #44754 
@jmg-duarte in the future, can you not force push your branch, but just push a commit that reflects your latest changes?  Thanks
I'm unsure what might have caused `test-arm` to fail as it didn't before
@jmg-duarte can you also look at https://github.com/pandas-dev/pandas/pull/44722/files#r763388551 .  Idea is to explicitly create those extension types in addition to the Decimal and Date ones
@jmg-duarte can you revert the 1.3.5 notes and fix conflicts on other, ping on green.
thanks @jmg-duarte very nice.

If you can followup with a test that deserializes a 0.20.0 format table schema would be great (it almost certainly works but would be great to have a specific test for this).
> thanks @jmg-duarte very nice.
> 
> If you can followup with a test that deserializes a 0.20.0 format table schema would be great (it almost certainly works but would be great to have a specific test for this).

I didn't mess with any previous schema tests so they work.
I can't explicitly test for 0.20.0 in the version because we upgraded it to 1.4.0.

Am I missing something?
@jmg-duarte thats the point we need an explicit example that provides backward compatibility 

make one in an older version of pandas and then save it in a file
Oh, you meant on reads.

A terça, 21 de dez de 2021, 13:06, Jeff Reback ***@***.***>
escreveu:

> @jmg-duarte <https://github.com/jmg-duarte> thats the point we need an
> explicit example that provides backward compatibility
>
> make one in an older version of pandas and then save it in a file
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/pandas-dev/pandas/pull/44722#issuecomment-998763154>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ADVCBSYJCX7FOOI3AC6WDQLUSB3UVANCNFSM5JHLVG2A>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>