diff --git a/openavmkit/cleaning.py b/openavmkit/cleaning.py index 479f8512..ae0382cf 100644 --- a/openavmkit/cleaning.py +++ b/openavmkit/cleaning.py @@ -745,8 +745,7 @@ def _fill_unknown_values(df, settings: dict): if cat_fields is not None: for field in cat_fields: if field in df: - df[field] = df[field].astype("str") - df[field] = df[field].fillna("UNKNOWN") + df[field] = df[field].astype("object").fillna("UNKNOWN").astype("str") if bool_fields is not None: for field in bool_fields: diff --git a/tests/test_cleaning.py b/tests/test_cleaning.py new file mode 100644 index 00000000..a56de6c0 --- /dev/null +++ b/tests/test_cleaning.py @@ -0,0 +1,23 @@ +import numpy as np +import pandas as pd + +from openavmkit.cleaning import _fill_unknown_values + + +def test_residual_categorical_fill_uses_unknown_not_nan(): + df = pd.DataFrame( + { + "key": ["a", "b", "c"], + "bldg_style": ["RAMBLER", np.nan, "SPLIT"], + } + ) + settings = { + "field_classification": { + "impr": {"categorical": ["bldg_style"]} + } + } + out = _fill_unknown_values(df, settings) + values = set(out["bldg_style"].astype(str)) + assert "UNKNOWN" in values + assert "nan" not in values + assert "" not in values