Fill missing categoricals as UNKNOWN instead of the string "nan"

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:

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 "<NA>" not in values