Repository: facultyai/marshmallow-dataframe

Issue:
Support categorical columns
Right now categorical columns are always recognised as a string dtype, which could be wrong. In addition, we are not validating that the deserialised values are one of the category options.

Discussion:

Patch:
diff --git a/src/marshmallow_dataframe/dataframe.py b/src/marshmallow_dataframe/dataframe.py
index 66369c5..7e612a9 100644
--- a/src/marshmallow_dataframe/dataframe.py
+++ b/src/marshmallow_dataframe/dataframe.py
@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import pandas as pd
 import numpy as np
 import marshmallow as ma
@@ -5,18 +7,29 @@ from typing import NamedTuple, List, Union, Dict
 
 __all__ = ["Dtypes", "RecordsDataFrameSchema", "SplitDataFrameSchema"]
 
-_FIELD_OPTIONS = {"required": True, "allow_none": True}
-
-# Integer columns in pandas cannot have null values, so here we allow_none for
-# all types except int
-DTYPE_KIND_TO_FIELD = {
-    "i": ma.fields.Int(required=True),
-    "u": ma.fields.Int(**_FIELD_OPTIONS),
-    "f": ma.fields.Float(allow_nan=True, **_FIELD_OPTIONS),
-    "O": ma.fields.Str(**_FIELD_OPTIONS),
-    "b": ma.fields.Bool(**_FIELD_OPTIONS),
-    "M": ma.fields.DateTime(**_FIELD_OPTIONS),
-    "m": ma.fields.TimeDelta(**_FIELD_OPTIONS),
+DTYPE_KIND_TO_FIELD_CLASS = {
+    "i": ma.fields.Int,
+    "u": ma.fields.Int,
+    "f": ma.fields.Float,
+    "O": ma.fields.Str,
+    "U": ma.fields.Str,
+    "S": ma.fields.Str,
+    "b": ma.fields.Bool,
+    "M": ma.fields.DateTime,
+    "m": ma.fields.TimeDelta,
+}
+
+_DEFAULT_FIELD_OPTIONS = {"required": True, "allow_none": True}
+
+DTYPE_KIND_TO_FIELD_OPTIONS: Dict[str, Dict[str, bool]] = defaultdict(
+    lambda: _DEFAULT_FIELD_OPTIONS
+)
+# Integer columns in pandas cannot have null values, so we allow_none for all
+# types except int
+DTYPE_KIND_TO_FIELD_OPTIONS["i"] = {"required": True}
+DTYPE_KIND_TO_FIELD_OPTIONS["f"] = {
+    "allow_nan": True,
+    **_DEFAULT_FIELD_OPTIONS,
 }
 
 
@@ -25,7 +38,7 @@ class Dtypes(NamedTuple):
     dtypes: List[np.dtype]
 
     @classmethod
-    def from_pandas_dtypes(cls, pd_dtypes: pd.DataFrame) -> "Dtypes":
+    def from_pandas_dtypes(cls, pd_dtypes: pd.Series) -> "Dtypes":
         return cls(
             columns=list(pd_dtypes.index), dtypes=list(pd_dtypes.values)
         )
@@ -39,6 +52,16 @@ class DtypeToFieldConversionError(Exception):
 
 
 def _dtype_to_field(dtype: np.dtype) -> ma.fields.Field:
+    # Object dtypes require more detailed mapping
+    if pd.api.types.is_categorical_dtype(dtype):
+        categories = dtype.categories.values.tolist()
+        kind = dtype.categories.dtype.kind
+        field_class = DTYPE_KIND_TO_FIELD_CLASS[kind]
+        field_options = DTYPE_KIND_TO_FIELD_OPTIONS[kind]
+        return field_class(
+            validate=ma.validate.OneOf(categories), **field_options
+        )
+
     try:
         kind = dtype.kind
     except AttributeError as exc:
@@ -48,12 +71,14 @@ def _dtype_to_field(dtype: np.dtype) -> ma.fields.Field:
         ) from exc
 
     try:
-        return DTYPE_KIND_TO_FIELD[kind]
+        field_class = DTYPE_KIND_TO_FIELD_CLASS[kind]
+        field_options = DTYPE_KIND_TO_FIELD_OPTIONS[kind]
+        return field_class(**field_options)
     except KeyError as exc:
         raise DtypeToFieldConversionError(
             f"The conversion of the dtype {dtype} with kind {dtype.kind} "
             "into marshmallow fields is unknown. Known kinds are: "
-            f"{DTYPE_KIND_TO_FIELD.keys()}"
+            f"{DTYPE_KIND_TO_FIELD_CLASS.keys()}"
         ) from exc
 
 
@@ -205,14 +230,7 @@ class RecordsDataFrameSchema(ma.Schema, metaclass=RecordsDataFrameSchemaMeta):
         index_data = {i: row for i, row in enumerate(records_data)}
         return pd.DataFrame.from_dict(
             index_data, orient="index", columns=self.opts.dtypes.columns
-        ).astype(
-            {
-                k: v
-                for k, v in zip(
-                    self.opts.dtypes.columns, self.opts.dtypes.dtypes
-                )
-            }
-        )
+        ).astype(dict(zip(self.opts.dtypes.columns, self.opts.dtypes.dtypes)))
 
 
 class SplitDataFrameSchema(ma.Schema, metaclass=SplitDataFrameSchemaMeta):
@@ -229,4 +247,8 @@ class SplitDataFrameSchema(ma.Schema, metaclass=SplitDataFrameSchemaMeta):
 
     @ma.post_load
     def make_df(self, data: dict) -> pd.DataFrame:
-        return pd.DataFrame(dtype=None, **data)
+        df = pd.DataFrame(dtype=None, **data).astype(
+            dict(zip(self.opts.dtypes.columns, self.opts.dtypes.dtypes))
+        )
+
+        return df


Test patch:
diff --git a/tests/test_dataframe.py b/tests/test_dataframe.py
index a87f974..3d28c9f 100644
--- a/tests/test_dataframe.py
+++ b/tests/test_dataframe.py
@@ -23,9 +23,14 @@ def sample_df():
             "string": ["a", "b"],
             "float": [1.1, 2.2],
             "datetime": ["2018-01-01T12:00:00Z", "2018-01-02T12:00:00Z"],
+            "categorical_int": [0, 1],
+            "categorical_str": ["zero", "one"],
         }
     )
 
+    df["categorical_int"] = df["categorical_int"].astype("category")
+    df["categorical_str"] = df["categorical_str"].astype("category")
+
     df["datetime"] = pd.to_datetime(df["datetime"]).astype("datetime64[ns]")
 
     return df
@@ -252,6 +257,31 @@ def test_records_schema_missing_data_field():
     )
 
 
+@pytest.mark.parametrize(
+    "column,data",
+    (["categorical_int", [0, 10]], ["categorical_str", ["one", "ten"]]),
+)
+def test_records_schema_wrong_category(sample_df, column, data):
+    class MySchema(RecordsDataFrameSchema):
+        class Meta:
+            dtypes = sample_df.dtypes
+
+    schema = MySchema()
+
+    # Replace categorical column data
+    new_df = sample_df.copy()
+    new_df[column] = data
+
+    input_df = serialize_df(new_df, orient="records")
+
+    with pytest.raises(ValidationError) as exception:
+        schema.load(input_df)
+
+    assert exception.value.messages["data"][1][column][0].startswith(
+        "Must be one of"
+    )
+
+
 @pytest.fixture
 def split_sample_schema(sample_df):
     class MySchema(SplitDataFrameSchema):
@@ -429,3 +459,26 @@ def test_split_schema_index_data_length_mismatch(
         exc.value.messages["data"][0]
         == f"Length of `index` and `data` must be equal."
     )
+
+
+@pytest.mark.parametrize(
+    "column,data",
+    (["categorical_int", [0, 10]], ["categorical_str", ["one", "ten"]]),
+)
+def test_split_schema_wrong_category(
+    split_sample_schema, sample_df, column, data
+):
+    # Replace categorical column data
+    new_df = sample_df.copy()
+    new_df[column] = data
+
+    column_index = list(new_df.columns).index(column)
+
+    input_df = serialize_df(new_df, orient="split")
+
+    with pytest.raises(ValidationError) as exception:
+        split_sample_schema.load(input_df)
+
+    assert exception.value.messages["data"][1][column_index][0].startswith(
+        "Must be one of"
+    )