fix(preprocessor): handle pandas categorical missing values

pretab/compose/factory.py

@@ -17,7 +17,7 @@
 
 from ..exceptions import ConfigWarning, IncompatibleParamsError, invalid_param_error
 from ..preprocessing.floats import ToFloatTransformer
-from ..preprocessing.missing import MissingStateIndicator
+from ..preprocessing.missing import MissingStateIndicator, _CategoricalMissingNormalizer
 from .config import PreprocessorConfig
 from .registry import (
     CATEGORICAL_ALIASES,
@@ -194,6 +194,8 @@ def get_categorical_transformer_steps(
 
     if add_imputer:
         imputer_kwargs = imputer_kwargs or {}
+        if "missing_values" not in imputer_kwargs:
+            steps.append(("normalize_missing", _CategoricalMissingNormalizer()))
         steps.append(
             ("imputer", SimpleImputer(strategy=imputer_strategy, add_indicator=add_missing_indicator, **imputer_kwargs))
         )
@@ -216,6 +218,17 @@ def get_categorical_transformer_steps(
         # Default to ignoring unseen categories so transform never crashes on
         # categories absent at fit time; callers can override via kwargs.
         onehot_kwargs = {"handle_unknown": "ignore", **kwargs}
+        categories = onehot_kwargs.get("categories", "auto")
+        drop = onehot_kwargs.get("drop")
+        if (
+            not add_imputer
+            and isinstance(categories, str)
+            and categories == "auto"
+            and (drop is None or isinstance(drop, str))
+        ):
+            # Explicit categories and dropped values keep sklearn's marker and
+            # numeric validation contracts; normalize only automatic encoding.
+            steps.append(("normalize_missing", _CategoricalMissingNormalizer()))
         steps.append(("onehot", cls(**onehot_kwargs)))
         steps.append(("to_float", ToFloatTransformer()))
     elif method == "pretrained":

pretab/compose/output.py

@@ -12,6 +12,7 @@
 """
 
 import numpy as np
+import pandas as pd
 from scipy import sparse as sp
 
 from ..exceptions import IncompatibleParamsError, OptionalDependencyError, PretabDataError
@@ -84,7 +85,14 @@ def compute_output_report(array, output_format, *, threshold=_SPARSE_AUTO_THRESH
         ``memory_saved_bytes``.
     """
     size = int(np.prod(array.shape, dtype=np.int64))
-    nonzero = int(array.count_nonzero()) if sp.issparse(array) else int(np.count_nonzero(array))
+    if sp.issparse(array):
+        nonzero = int(array.count_nonzero())
+    elif array.dtype.kind == "O":
+        # pd.NA has no boolean value; count it like NaN without modifying the
+        # raw output or the truthiness of any other object value.
+        nonzero = sum(value is pd.NA or bool(value) for value in array.flat)
+    else:
+        nonzero = int(np.count_nonzero(array))
     density = float(nonzero) / size if size else 0.0
     dense_bytes = size * int(array.dtype.itemsize)
 

pretab/encoding/categorical/ordinal.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils.validation import check_is_fitted
 
@@ -7,20 +8,15 @@
 
 
 def _is_missing_value(value) -> bool:
-    """Return True for ``None`` or a NaN float; False for any other value."""
-    if value is None:
-        return True
-    try:
-        return bool(np.isnan(value))
-    except TypeError:
-        return False
+    """Return whether a scalar category is missing according to pandas."""
+    return bool(pd.isna(value))
 
 
 class ContinuousOrdinalTransformer(RepresentationSpecMixin, TransformerMixin, BaseEstimator):
     """Encode categorical features as continuous integer values.
 
     Each unique category within a feature is assigned an integer based on its
-    sorted order. Unknown, missing (``None`` or NaN), or unseen categories are
+    sorted order. Unknown, missing (``None``, ``pd.NA`` or NaN), or unseen categories are
     mapped to ``0``. This is useful for models that can only handle numerical
     input.
 
@@ -32,7 +28,7 @@ class ContinuousOrdinalTransformer(RepresentationSpecMixin, TransformerMixin, Ba
     Notes
     -----
     Categories are numbered starting at ``1`` in sorted order; the value ``0`` is
-    reserved for categories not seen during ``fit`` (and for ``None`` / NaN).
+    reserved for categories not seen during ``fit`` (and for ``None`` / ``pd.NA`` / NaN).
 
     Examples
     --------
@@ -69,7 +65,7 @@ def fit(self, X, y=None):
         self.mapping_ = []
         for j in range(X.shape[1]):
             column = X[:, j]
-            # Missing values (None / NaN) are excluded before sorting: np.unique
+            # Missing values are excluded before sorting: np.unique
             # cannot compare a NaN or None against a string category.
             non_missing = np.array([v for v in column if not _is_missing_value(v)], dtype=object)
             categories = np.unique(non_missing) if non_missing.size else np.array([], dtype=object)
@@ -104,7 +100,7 @@ def transform(self, X):
             )
         out = np.zeros(X.shape, dtype=int)
         for j, mapping in enumerate(self.mapping_):
-            out[:, j] = [mapping.get(v, 0) for v in X[:, j]]
+            out[:, j] = [0 if _is_missing_value(v) else mapping.get(v, 0) for v in X[:, j]]
         return out
 
     def get_feature_names_out(self, input_features=None):

pretab/preprocessing/missing.py

@@ -1,7 +1,35 @@
 import numpy as np
 import pandas as pd
-from sklearn.base import BaseEstimator, TransformerMixin
-from sklearn.utils.validation import check_is_fitted
+from scipy import sparse as sp
+from sklearn.base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
+from sklearn.utils.validation import check_array, check_is_fitted, validate_data
+
+
+class _CategoricalMissingNormalizer(OneToOneFeatureMixin, TransformerMixin, BaseEstimator):
+    """Normalize object missing markers without changing numeric or sparse inputs."""
+
+    def fit(self, X, y=None):
+        validate_data(self, X, skip_check_array=True)
+        return self
+
+    def transform(self, X):
+        check_is_fitted(self, "n_features_in_")
+        validate_data(self, X, reset=False, skip_check_array=True)
+        if sp.issparse(X) or (
+            isinstance(X, pd.DataFrame) and all(isinstance(dtype, pd.SparseDtype) for dtype in X.dtypes)
+        ):
+            # Keep pandas sparse columns intact: converting them to scipy sparse
+            # rejects nonzero fill values and changes OneHotEncoder's input contract.
+            return X.copy()
+        values = check_array(X, dtype=None, ensure_all_finite=False, copy=True)  # type: ignore[arg-type]
+        if values.dtype.kind == "O":
+            return np.where(pd.isna(values), np.nan, values)
+        return X.copy() if isinstance(X, pd.DataFrame) else values
+
+    def __sklearn_tags__(self):
+        tags = super().__sklearn_tags__()  # type: ignore[attr-defined]
+        tags.input_tags.allow_nan = True
+        return tags
 
 
 class MissingStateIndicator(TransformerMixin, BaseEstimator):

tests/compose/test_categorical_pipeline.py

@@ -1,5 +1,7 @@
 import numpy as np
+import pandas as pd
 import pytest
+from scipy import sparse as sp
 from sklearn.pipeline import Pipeline
 
 from pretab.compose.factory import get_categorical_transformer_steps
@@ -9,6 +11,12 @@ def _build(method, **kwargs):
     return Pipeline(get_categorical_transformer_steps(method, add_imputer=False, **kwargs))
 
 
+def _feature_names(pipe):
+    names = pipe.get_feature_names_out()
+    assert names is not None
+    return names.tolist()
+
+
 def test_one_hot_ignores_unseen_categories():
     pipe = _build("one-hot", sparse_output=False)
     pipe.fit(np.array([["A"], ["B"], ["C"]]))
@@ -27,3 +35,109 @@ def test_one_hot_handle_unknown_override():
     pipe.fit(np.array([["A"], ["B"]]))
     with pytest.raises(ValueError):
         pipe.transform(np.array([["C"]]))
+
+
+@pytest.mark.parametrize("marker", [None, pd.NA, np.nan], ids=["None", "pd.NA", "NaN"])
+def test_one_hot_missing_after_clean_fit_matches_unseen(marker):
+    pipe = _build("one-hot", sparse_output=False).fit(pd.DataFrame({"c": ["a", "b"]}))
+
+    result = pipe.transform(pd.DataFrame({"c": [marker, "unseen"]}))
+
+    np.testing.assert_array_equal(result, np.zeros((2, 2)))
+
+
+@pytest.mark.parametrize("marker", [None, pd.NA, np.nan], ids=["None", "pd.NA", "NaN"])
+def test_one_hot_all_missing_matches_nan(marker):
+    frame = pd.DataFrame({"c": [marker, marker, marker]}, dtype=object)
+    reference = pd.DataFrame({"c": [np.nan, np.nan, np.nan]}, dtype=object)
+    expected = _build("one-hot", sparse_output=False).fit(reference)
+    actual = _build("one-hot", sparse_output=False).fit(frame)
+
+    assert _feature_names(actual) == _feature_names(expected)
+    np.testing.assert_array_equal(actual.transform(frame), expected.transform(reference))
+
+
+def test_one_hot_explicit_none_category_is_preserved():
+    categories = [["a", "b", None]]
+    pipe = _build("one-hot", categories=categories, sparse_output=False)
+    frame = pd.DataFrame({"c": ["a", None, "b"]})
+
+    result = pipe.fit_transform(frame)
+
+    np.testing.assert_array_equal(result, [[1, 0, 0], [0, 0, 1], [0, 1, 0]])
+    assert _feature_names(pipe) == ["c_a", "c_b", "c_None"]
+    assert categories == [["a", "b", None]]
+
+
+def test_one_hot_explicit_none_drop_is_preserved():
+    drop = [None]
+    pipe = _build("one-hot", drop=drop, sparse_output=False)
+
+    result = pipe.fit_transform(pd.DataFrame({"c": ["a", None, "b"]}))
+
+    np.testing.assert_array_equal(result, [[1, 0], [0, 0], [0, 1]])
+    assert _feature_names(pipe) == ["c_a", "c_b"]
+    assert drop == [None]
+
+
+@pytest.mark.parametrize("add_imputer", [True, False])
+@pytest.mark.parametrize("nullable", [True, False])
+def test_one_hot_unsorted_numeric_categories_are_rejected(add_imputer, nullable):
+    frame = pd.DataFrame({"c": pd.Series([1, 2], dtype="Int64")}) if nullable else np.array([[1], [2]])
+    pipe = Pipeline(
+        get_categorical_transformer_steps("one-hot", add_imputer=add_imputer, categories=[[2, 1]], sparse_output=False)
+    )
+
+    with pytest.raises(ValueError, match="Unsorted categories"):
+        pipe.fit(frame)
+
+
+@pytest.mark.parametrize("add_imputer", [True, False])
+def test_one_hot_clean_numeric_fit_with_missing_at_transform(add_imputer):
+    reference = pd.DataFrame({"c": [1.0, 2.0]})
+    frame = pd.DataFrame({"c": [1, 2]})
+    options = {"add_imputer": add_imputer, "sparse_output": False}
+    expected = Pipeline(get_categorical_transformer_steps("one-hot", **options)).fit(reference)
+    actual = Pipeline(get_categorical_transformer_steps("one-hot", **options)).fit(frame)
+
+    np.testing.assert_array_equal(
+        actual.transform(pd.DataFrame({"c": [None, 2]})),
+        expected.transform(pd.DataFrame({"c": [np.nan, 2.0]})),
+    )
+
+
+def test_one_hot_pandas_sparse_columns_keep_native_support():
+    frame = pd.DataFrame({"c": pd.arrays.SparseArray([1, 0, 2], fill_value=0)})
+    original = frame.copy(deep=True)
+    dense = pd.DataFrame({"c": [1, 0, 2]})
+    actual = _build("one-hot", sparse_output=False).fit(frame)
+    expected = _build("one-hot", sparse_output=False).fit(dense)
+
+    np.testing.assert_array_equal(actual.transform(frame), expected.transform(dense))
+    assert _feature_names(actual) == _feature_names(expected)
+    pd.testing.assert_frame_equal(frame, original)
+
+
+@pytest.mark.parametrize("fill_value", [0.0, 1.0, np.nan], ids=["zero", "nonzero", "NaN"])
+def test_one_hot_pandas_sparse_missing_keeps_native_support(fill_value):
+    frame = pd.DataFrame({"c": pd.arrays.SparseArray([1.0, np.nan, 2.0], fill_value=fill_value)})
+    original = frame.copy(deep=True)
+    dense = pd.DataFrame({"c": [1.0, np.nan, 2.0]})
+    actual = _build("one-hot", sparse_output=False).fit(frame)
+    expected = _build("one-hot", sparse_output=False).fit(dense)
+
+    np.testing.assert_array_equal(actual.transform(frame), expected.transform(dense))
+    assert _feature_names(actual) == _feature_names(expected)
+    pd.testing.assert_frame_equal(frame, original)
+
+
+def test_categorical_imputation_preserves_sparse_input_support():
+    frame = sp.csr_matrix([[1.0], [np.nan], [2.0]])
+    original = frame.copy()
+    pipe = Pipeline(get_categorical_transformer_steps("none"))
+
+    result = pipe.fit_transform(frame)
+
+    assert sp.issparse(result)
+    np.testing.assert_array_equal(result.toarray(), [[1.0], [1.0], [2.0]])
+    np.testing.assert_array_equal(frame.toarray(), original.toarray())

tests/compose/test_factory.py

@@ -1,8 +1,10 @@
 """Unit tests for :mod:`pretab.compose.factory`."""
 
 import numpy as np
+import pandas as pd
 import pytest
 from sklearn.compose import ColumnTransformer
+from sklearn.pipeline import Pipeline
 
 from pretab.compose.factory import (
     _placement_kwargs,
@@ -75,7 +77,9 @@ def test_unknown_numerical_method_raises():
 # categorical step assembly
 # --------------------------------------------------------------------------- #
 def test_one_hot_appends_to_float():
-    assert _names(get_categorical_transformer_steps("one-hot", add_imputer=False)) == ["onehot", "to_float"]
+    pipe = Pipeline(get_categorical_transformer_steps("one-hot", add_imputer=False))
+    transformed = pipe.fit_transform(np.array([["a"], ["b"]]))
+    assert transformed.dtype == np.float64
 
 
 def test_int_uses_continuous_ordinal():
@@ -87,6 +91,38 @@ def test_unknown_categorical_method_raises():
         get_categorical_transformer_steps("does-not-exist", add_imputer=False)
 
 
+def test_explicit_numerical_missing_value_override_is_preserved():
+    options = {"missing_values": -1}
+    pipe = Pipeline(get_numerical_transformer_steps("none", imputer_kwargs=options))
+    result = pipe.fit_transform(np.array([[1.0], [-1.0], [3.0]]))
+
+    np.testing.assert_array_equal(result.ravel(), [1.0, 2.0, 3.0])
+    assert options == {"missing_values": -1}
+
+
+def test_explicit_categorical_missing_value_override_is_preserved():
+    options = {"missing_values": "missing"}
+    pipe = Pipeline(get_categorical_transformer_steps("int", imputer_kwargs=options))
+    result = pipe.fit_transform(np.array([["a"], ["missing"], ["b"]], dtype=object))
+
+    np.testing.assert_array_equal(result.ravel(), [1, 1, 2])
+    assert options == {"missing_values": "missing"}
+
+
+@pytest.mark.parametrize("missing_values", [None, np.nan, pd.NA], ids=["None", "NaN", "pd.NA"])
+def test_explicit_imputer_missing_marker_is_authoritative(missing_values):
+    options = {"missing_values": missing_values}
+    pipe = Pipeline(get_categorical_transformer_steps("none", imputer_kwargs=options))
+
+    result = pipe.fit_transform(np.array([["a"], [None], ["b"], ["a"]], dtype=object))
+
+    if missing_values is np.nan:
+        assert result[1, 0] is None
+    else:
+        assert result[1, 0] == "a"
+    assert options["missing_values"] is missing_values
+
+
 # --------------------------------------------------------------------------- #
 # placement kwargs by capability class
 # --------------------------------------------------------------------------- #

tests/encoding/categorical/test_continuous_ordinal_transformer.py

@@ -48,3 +48,21 @@ def test_continuous_ordinal_unseen_category_maps_to_zero():
     transformer = ContinuousOrdinalTransformer().fit(np.array([["a"], ["b"]], dtype=object))
     Xt = transformer.transform(np.array([["c"]], dtype=object))
     assert Xt.ravel().tolist() == [0]
+
+
+def test_continuous_ordinal_handles_nullable_pandas_strings():
+    frame = pd.DataFrame({"c": ["red", None, "blue", "red"]}).convert_dtypes()
+
+    transformed = ContinuousOrdinalTransformer().fit_transform(frame)
+
+    assert transformed.ravel().tolist() == [2, 0, 1, 2]
+
+
+def test_continuous_ordinal_all_missing_and_transform_time_missing():
+    missing = pd.DataFrame({"c": [None, pd.NA, np.nan]}, dtype=object)
+    all_missing = ContinuousOrdinalTransformer().fit(missing)
+    np.testing.assert_array_equal(all_missing.transform(missing), np.zeros((3, 1), dtype=int))
+
+    clean = ContinuousOrdinalTransformer().fit(pd.DataFrame({"c": ["a", "b"]}))
+    transformed = clean.transform(pd.DataFrame({"c": [pd.NA, "unseen", "a"]}))
+    assert transformed.ravel().tolist() == [0, 0, 1]

tests/integration/test_output_format.py

@@ -166,6 +166,26 @@ def test_output_report_shape_and_keys(frame, y):
     assert 0.0 <= report["density"] <= 1.0
 
 
+@pytest.mark.parametrize("return_array", [True, False])
+def test_object_output_density_keeps_raw_missing_markers(return_array):
+    frame = pd.DataFrame({"c": ["a", pd.NA, None, np.nan, 0, False, ""]}, dtype=object)
+    original = frame.copy(deep=True)
+    fitted = Preprocessor(categorical_method="none", missing_policy="propagate", output_format="dense").fit(frame)
+
+    transformed = fitted.transform(frame, return_array=return_array)
+    result = transformed if return_array else transformed["cat_c"]
+
+    assert fitted.output_report_["density"] == pytest.approx(3 / 7)
+    assert result[0, 0] == "a"
+    assert result[1, 0] is pd.NA
+    assert result[2, 0] is None
+    assert np.isnan(result[3, 0])
+    assert result[4, 0] == 0
+    assert result[5, 0] is False
+    assert result[6, 0] == ""
+    pd.testing.assert_frame_equal(frame, original)
+
+
 # --- set_output ----------------------------------------------------------------