From 964f5d31774fd887be1e91086824ab513fcd5155 Mon Sep 17 00:00:00 2001 From: Nikhil Singhal Date: Sat, 4 Jul 2026 11:52:50 -0700 Subject: [PATCH] Residual categorical fill: fill nulls before the string cast astype("str") converts NaN to the literal string "nan" before fillna runs, so models received a "nan" category instead of "UNKNOWN". Fill on object dtype first, then cast. Adds tests/test_cleaning.py covering the null-to-UNKNOWN path. --- openavmkit/cleaning.py | 3 +-- tests/test_cleaning.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) create mode 100644 tests/test_cleaning.py 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