fix(preprocessor): make handle_missing="error" actually reject NaN

pretab/pipeline/numerical.py

@@ -4,6 +4,7 @@
 from sklearn.preprocessing import MinMaxScaler, StandardScaler
 
 from ..core.exceptions import ConfigWarning, invalid_param_error
+from ..transformers.encoders.floats import RaiseOnNaNTransformer
 from .registry import NUMERICAL_ALIASES, NUMERICAL_METHODS, resolve_method
 
 # Spline basis expansions that share the target-aware knot API.
@@ -121,6 +122,13 @@ def get_numerical_transformer_steps(
     if add_imputer:
         imputer_kwargs = imputer_kwargs or {}
         steps.append(("imputer", SimpleImputer(strategy=imputer_strategy, **imputer_kwargs)))
+    else:
+        # ``handle_missing="error"`` drops the imputer. Enforce the policy here so
+        # every method rejects NaN, rather than leaving it to whether the chosen
+        # transformer happens to notice: the scikit-learn scalers ignore missing
+        # values by design and the PreTab families declare ``allow_nan``, so only
+        # PLE used to raise and everything else emitted NaN silently.
+        steps.append(("nan_check", RaiseOnNaNTransformer()))
 
     # Define scalers that could be added independently
     scalers = {

pretab/transformers/__init__.py

@@ -3,6 +3,7 @@
 from .encoders import (
     ContinuousOrdinalTransformer,
     NoTransformer,
+    RaiseOnNaNTransformer,
     ToFloatTransformer,
 )
 from .feature_maps import (
@@ -45,6 +46,7 @@
     "PLETransformer",
     "PSplineTransformer",
     "RBFExpansionTransformer",
+    "RaiseOnNaNTransformer",
     "ReLUExpansionTransformer",
     "RollingStatsTransformer",
     "SigmoidExpansionTransformer",

pretab/transformers/encoders/__init__.py

@@ -5,10 +5,11 @@
 """
 
 from .continuous_ordinal import ContinuousOrdinalTransformer
-from .floats import NoTransformer, ToFloatTransformer
+from .floats import NoTransformer, RaiseOnNaNTransformer, ToFloatTransformer
 
 __all__ = [
     "ContinuousOrdinalTransformer",
     "NoTransformer",
+    "RaiseOnNaNTransformer",
     "ToFloatTransformer",
 ]

pretab/transformers/encoders/floats.py

@@ -2,6 +2,8 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils.validation import check_is_fitted
 
+from ...core.exceptions import PretabDataError
+
 
 class NoTransformer(TransformerMixin, BaseEstimator):
     """Pass-through transformer that returns the input unchanged.
@@ -160,3 +162,111 @@ def __sklearn_tags__(self):
         tags = super().__sklearn_tags__()
         tags.input_tags.allow_nan = True
         return tags
+
+
+class RaiseOnNaNTransformer(TransformerMixin, BaseEstimator):
+    """Pass data through, raising a clear error if it contains missing values.
+
+    Sits at the head of the numerical pipeline when
+    ``Preprocessor(handle_missing="error")`` drops the imputer, so the policy is
+    enforced in one place for every numerical method.
+
+    Without it the policy depended on whether the chosen transformer happened to
+    notice NaN: the plain scikit-learn scalers ignore missing values by design,
+    the PreTab expansion families declare ``allow_nan`` so they let them through,
+    and only ``PLETransformer`` actually raised. Everything else silently emitted
+    a NaN-contaminated feature matrix -- and the unsupervised feature-map path was
+    worse than pass-through, because ``np.percentile`` over a column containing
+    NaN makes *every* center NaN.
+
+    Attributes
+    ----------
+    n_features_in_ : int
+        Number of input features seen during ``fit``.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from pretab.transformers import RaiseOnNaNTransformer
+    >>> RaiseOnNaNTransformer().fit_transform(np.array([[1.0], [2.0]])).shape
+    (2, 1)
+    """
+
+    def fit(self, X, y=None):
+        """Check ``X`` for missing values and record the feature count.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            The input data to check.
+        y : Ignored
+            Not used, present for API consistency by convention.
+
+        Returns
+        -------
+        self : object
+            Fitted transformer.
+        """
+        X = self._check_no_nan(X)
+        self.n_features_in_ = X.shape[1]
+        return self
+
+    def transform(self, X):
+        """Return the input unchanged, raising if it contains missing values.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            The input data to check.
+
+        Returns
+        -------
+        X : ndarray
+            The validated input, as a float array.
+
+        Raises
+        ------
+        pretab.core.exceptions.PretabDataError
+            If ``X`` contains any NaN.
+        """
+        check_is_fitted(self, "n_features_in_")
+        return self._check_no_nan(X)
+
+    @staticmethod
+    def _check_no_nan(X) -> np.ndarray:
+        """Coerce to a 2D float array and reject missing values."""
+        array = np.asarray(X, dtype=np.float64)
+        if array.ndim == 1:
+            array = array.reshape(-1, 1)
+        if np.isnan(array).any():
+            raise PretabDataError(
+                "Input contains NaN, but handle_missing='error' was requested.\n"
+                "Fix: pass handle_missing='median' to impute missing values before "
+                "the numerical method, or remove them from the input."
+            )
+        return array
+
+    def get_feature_names_out(self, input_features=None):
+        """Return the output feature names (unchanged from the input).
+
+        Parameters
+        ----------
+        input_features : list of str or None
+            The names of the input features. When ``None``, names of the form
+            ``x0, x1, ...`` are generated.
+
+        Returns
+        -------
+        feature_names : ndarray of shape (n_features,)
+            The output feature names.
+        """
+        check_is_fitted(self, "n_features_in_")
+        if input_features is None:
+            input_features = [f"x{i}" for i in range(self.n_features_in_)]
+        return np.asarray(input_features, dtype=object)
+
+    def __sklearn_tags__(self):
+        """Declare that missing values are rejected, not passed through."""
+        tags = super().__sklearn_tags__()
+        tags.input_tags.allow_nan = False
+        return tags

tests/test_encoder_feature_counts.py

@@ -34,3 +34,40 @@ def test_custom_bin_transformer_reads_actual_column_count():
     # Proves the value is derived from X, not hardcoded to 1.
     X = np.zeros((10, 2))
     assert CustomBinTransformer(output_dim=4).fit(X).n_features_in_ == 2
+
+
+# --------------------------------------------------------------------------- #
+# RaiseOnNaNTransformer: the single enforcement point for handle_missing="error"
+# --------------------------------------------------------------------------- #
+def test_raise_on_nan_passes_clean_data_through():
+    from pretab.transformers import RaiseOnNaNTransformer
+
+    X = np.array([[1.0, 2.0], [3.0, 4.0]])
+    out = RaiseOnNaNTransformer().fit_transform(X)
+
+    np.testing.assert_array_equal(out, X)
+
+
+def test_raise_on_nan_rejects_missing_values_at_fit():
+    from pretab.core.exceptions import PretabDataError
+    from pretab.transformers import RaiseOnNaNTransformer
+
+    with pytest.raises(PretabDataError, match="handle_missing"):
+        RaiseOnNaNTransformer().fit(np.array([[1.0], [np.nan]]))
+
+
+def test_raise_on_nan_rejects_missing_values_at_transform():
+    from pretab.core.exceptions import PretabDataError
+    from pretab.transformers import RaiseOnNaNTransformer
+
+    transformer = RaiseOnNaNTransformer().fit(np.array([[1.0], [2.0]]))
+
+    with pytest.raises(PretabDataError, match="handle_missing"):
+        transformer.transform(np.array([[1.0], [np.nan]]))
+
+
+@pytest.mark.parametrize("n_cols", [1, 2, 3])
+def test_raise_on_nan_records_feature_count(n_cols):
+    from pretab.transformers import RaiseOnNaNTransformer
+
+    assert RaiseOnNaNTransformer().fit(np.zeros((5, n_cols))).n_features_in_ == n_cols

tests/test_reproducibility.py

@@ -11,6 +11,8 @@
 import pytest
 from sklearn.base import clone
 
+from pretab.core.exceptions import PretabDataError
+from pretab.pipeline import get_numerical_transformer_steps
 from pretab.preprocessor import Preprocessor
 from pretab.transformers import RBFExpansionTransformer
 
@@ -118,6 +120,67 @@ def test_handle_missing_error_rejects_nan(data):
         pre.fit(X, y)
 
 
+# The "error" policy must hold for *every* numerical method, not just PLE.
+#
+# ``handle_missing`` only ever dropped the imputer and was forwarded to PLE
+# alone, so the guarantee depended on whether the chosen transformer happened to
+# notice NaN. The scikit-learn scalers ignore missing values by design and the
+# PreTab families declare ``allow_nan``, so everything except PLE silently
+# emitted a NaN-contaminated matrix -- and the unsupervised feature maps were
+# worse still, because ``np.percentile`` over a NaN column makes every center NaN.
+@pytest.mark.parametrize(
+    "method",
+    ["minmax", "standardization", "robust", "quantile", "rbf", "relu", "tanh",
+     "cubicspline", "pspline", "tprs", "none", "ple"],
+)
+def test_handle_missing_error_rejects_nan_for_every_method(data, method):
+    X, y = data
+    X = X.copy()
+    X.iloc[0, 0] = np.nan
+
+    with pytest.raises(ValueError):
+        Preprocessor(numerical_method=method, handle_missing="error").fit(X, y)
+
+
+@pytest.mark.parametrize("method", ["minmax", "rbf", "cubicspline"])
+def test_handle_missing_median_still_imputes_for_every_method(data, method):
+    X, y = data
+    X = X.copy()
+    X.iloc[0, 0] = np.nan
+
+    out = Preprocessor(numerical_method=method, handle_missing="median").fit(X, y).transform(
+        X, return_array=True
+    )
+    assert isinstance(out, np.ndarray)
+    assert np.isfinite(out).all()
+
+
+def test_handle_missing_error_raises_a_pretab_error_naming_the_option(data):
+    X, y = data
+    X = X.copy()
+    X.iloc[0, 0] = np.nan
+
+    with pytest.raises(PretabDataError, match="handle_missing='error'"):
+        Preprocessor(numerical_method="minmax", handle_missing="error").fit(X, y)
+
+
+def test_nan_check_step_only_present_when_erroring():
+    erroring = [name for name, _ in get_numerical_transformer_steps("minmax", add_imputer=False)]
+    imputing = [name for name, _ in get_numerical_transformer_steps("minmax", add_imputer=True)]
+
+    assert "nan_check" in erroring and "imputer" not in erroring
+    assert "imputer" in imputing and "nan_check" not in imputing
+
+
+def test_handle_missing_error_still_transforms_clean_data(data):
+    X, y = data
+    out = Preprocessor(numerical_method="rbf", handle_missing="error").fit(X, y).transform(
+        X, return_array=True
+    )
+    assert isinstance(out, np.ndarray)
+    assert np.isfinite(out).all()
+
+
 # --- transformer / helper level seeding ------------------------------------ #
 
 def test_rbf_transformer_seeded_centers_reproducible(data):