fix(core): apply the all-NaN column fill to training data only

README.md

@@ -276,7 +276,7 @@ model.fit(X_train, y_train, max_epochs=50)
 
 > **Learn more:** Preprocessing is driven by `PreprocessingConfig`; see the [Config System](https://deeptab.readthedocs.io/en/latest/core_concepts/config_system.html) guide and the [PreTab](https://github.com/OpenTabular/PreTab) project.
 
-> **Note:** Duplicate column names raise a `DuplicateColumnsError`, and a column that is entirely missing is imputed with a constant rather than passed through untouched. See the [FAQ](https://deeptab.readthedocs.io/en/latest/getting_started/faq.html#data-and-preprocessing) for details.
+> **Note:** Duplicate column names raise a `DuplicateColumnsError`, and a training column that is entirely missing is imputed with a constant rather than passed through untouched. See the [FAQ](https://deeptab.readthedocs.io/en/latest/getting_started/faq.html#data-and-preprocessing) for details.
 
 ### Observability & Experiment Tracking
 

deeptab/core/inference.py

@@ -322,7 +322,7 @@ def validate_input(
         >>> X_valid = model.validate_input(X_new)
         >>> predictions = model.predict(X_valid)
         """
-        X_df = ensure_dataframe(X)
+        X_df = ensure_dataframe(X, context="predict")
 
         expected_names = self.feature_names
 

deeptab/core/sklearn_compat.py

@@ -36,16 +36,20 @@ def ensure_dataframe(X: Any, context: str = "fit") -> pd.DataFrame:
       sentinel the input used.
     - Any remaining non-numeric, non-object column dtype raises
       :exc:`~deeptab.core.exceptions.ColumnDtypeError` naming each offending column.
-    - Columns where every value is NaN are filled with a constant (0 for numeric
-      columns, a sentinel category for object columns) and issue a
-      :class:`~deeptab.core.exceptions.DataWarning`.
+    - For training data (``context="fit"``), columns where every value is NaN are
+      filled with a constant (0 for numeric columns, a sentinel category for object
+      columns) and issue a :class:`~deeptab.core.exceptions.DataWarning`. Validation
+      and prediction inputs keep their NaN values so the imputer fitted on the
+      training data fills them; a single-row prediction with a missing value would
+      otherwise count as an all-NaN column.
 
     Parameters
     ----------
     X:
         Input data.  Converted to :class:`pandas.DataFrame` if necessary.
     context:
-        Name of the calling method (used in error messages).
+        Name of the calling method (used in error messages). Only ``"fit"``
+        applies the all-NaN column fill.
     """
     # Reject 1-D input early: sklearn convention requires 2-D feature arrays.
     _arr = np.asarray(X) if not isinstance(X, pd.DataFrame | pd.Series) else X
@@ -112,7 +116,7 @@ def ensure_dataframe(X: Any, context: str = "fit") -> pd.DataFrame:
     # them with a constant here rather than leaving it to PreTab's own imputer, which
     # cannot compute a mean/median from zero observed values and simply drops the
     # column instead.
-    all_nan_cols = [c for c in df.columns if bool(df[c].isna().all())]
+    all_nan_cols = [c for c in df.columns if bool(df[c].isna().all())] if context == "fit" else []
     if all_nan_cols:
         df = df.copy()
         for c in all_nan_cols:

deeptab/models/_mixins/fit.py

@@ -156,7 +156,7 @@ def _build_model(
         if hasattr(y, "values"):
             y = y.values
         if X_val is not None:
-            X_val = ensure_dataframe(X_val)
+            X_val = ensure_dataframe(X_val, context="validation")
             if y_val is not None and hasattr(y_val, "values"):
                 y_val = y_val.values
 

deeptab/models/lss_base.py

@@ -176,7 +176,7 @@ def build_model(
         if hasattr(y, "values"):
             y = y.values
         if X_val is not None:
-            X_val = ensure_dataframe(X_val)
+            X_val = ensure_dataframe(X_val, context="validation")
             if y_val is not None and hasattr(y_val, "values"):
                 y_val = y_val.values
         if self.family_name == "categorical":

docs/getting_started/faq.md

@@ -89,7 +89,7 @@ model.fit(df, y, max_epochs=50)
 The internal [PreTab](https://github.com/OpenTabular/PreTab) preprocessor imputes missing values as part of fitting, so you do not need a separate imputation step. The exact strategy follows the configured `PreprocessingConfig`; with the defaults it uses PreTab's built-in imputation for numerical and categorical features.
 
 ```{note}
-A column that is entirely missing (every row is `NaN`/`None`) has nothing for PreTab to learn a strategy from, so DeepTab fills it with a constant instead (`0` for numeric columns, `"missing"` for object/categorical columns) and emits a `DataWarning`. Consider dropping such columns before calling `fit()`.
+A training column that is entirely missing (every row is `NaN`/`None`) has nothing for PreTab to learn a strategy from, so DeepTab fills it with a constant instead (`0` for numeric columns, `"missing"` for object/categorical columns) and emits a `DataWarning`. Consider dropping such columns before calling `fit()`. Validation and prediction inputs are never filled this way: their missing values, including a single-row request with a missing value, are imputed with the strategy learned during `fit()`.
 ```
 
 ### What if my column names have duplicates?

tests/core/test_exceptions.py

@@ -1133,6 +1133,25 @@ def test_all_nan_object_column_filled_with_sentinel(self):
             result = ensure_dataframe(df)
         assert (result["all_nan"] == "missing").all()
 
+    @pytest.mark.parametrize("context", ["predict", "validation"])
+    def test_all_nan_column_is_left_to_fitted_imputer_outside_fit(self, context):
+        import warnings
+
+        from deeptab.core.sklearn_compat import ensure_dataframe
+
+        df = pd.DataFrame(
+            {
+                "good": [1.0],
+                "num": [np.nan],
+                "cat": pd.Series([None], dtype="object"),
+            }
+        )
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", DataWarning)
+            result = ensure_dataframe(df, context=context)
+        assert bool(result["num"].isna().all())
+        assert bool(result["cat"].isna().all())
+
     def test_duplicate_columns_raise_duplicate_columns_error(self):
         from deeptab.core.sklearn_compat import ensure_dataframe
 

tests/integration/test_data_workflows.py

@@ -136,3 +136,35 @@ def test_dataset_with_batch_object_mode(self, simple_tensors):
         # Test tuple conversion
         batch_tuple = batch.to_tuple()
         assert isinstance(batch_tuple, tuple)
+
+
+@pytest.mark.parametrize("missing_column", ["num", "cat"])
+def test_single_row_with_missing_value_matches_batched_prediction(missing_column):
+    """A missing value is imputed by the fitted preprocessor, whatever the batch size."""
+    import numpy as np
+
+    from deeptab.configs import TrainerConfig
+    from deeptab.core.inference import InferenceModel
+    from deeptab.models import MLPRegressor
+
+    rng = np.random.default_rng(0)
+    X = pd.DataFrame(
+        {
+            "num": rng.normal(50, 10, 120),
+            "cat": rng.choice(["a", "b", "c"], 120).astype(object),
+            "other": rng.normal(size=120),
+        }
+    )
+    X.loc[::10, "num"] = np.nan
+    X.loc[5::10, "cat"] = np.nan
+    y = X["other"].to_numpy() + rng.normal(size=120)
+    model = MLPRegressor(trainer_config=TrainerConfig(max_epochs=1, batch_size=32), random_state=0)
+    model.fit(X, y, accelerator="cpu", logger=False, enable_progress_bar=False, enable_model_summary=False)
+
+    row = X.iloc[[1]].copy()
+    row[missing_column] = np.nan
+    batch = pd.concat([row, X.iloc[[2]]], ignore_index=True)
+    np.testing.assert_allclose(model.predict(row), model.predict(batch)[:1], rtol=1e-5)
+
+    validated = InferenceModel.from_estimator(model).validate_input(row)
+    assert bool(validated[missing_column].isna().all())