Skip to content

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

Open
ChrisW09 wants to merge 2 commits into
mainfrom
fix/predict-nan-imputation
Open

ChrisW09 wants to merge 2 commits into
mainfrom
fix/predict-nan-imputation

Conversation

@ChrisW09

Copy link
Copy Markdown
Collaborator

Fixes #495

Problem

A missing value in a single-row prediction request was replaced by 0 (numeric) or "missing" (categorical) instead of being imputed by the preprocessor fitted on the training data. In the issue's example, a row with a missing age was predicted as 15.8 on its own (exactly the prediction for age=0) but as 49.7 inside a two-row batch.

Root cause

ensure_dataframe() fills entirely-NaN columns with a constant. That safeguard exists for training data: an all-NaN training column gives PreTab nothing to learn an imputation from. The fill also ran on prediction input (validate_input_features), on InferenceModel.validate_input() (which used the default context="fit") and on explicit X_val, where a single row with a missing value already forms an all-NaN column.

Fix

  • ensure_dataframe() applies the all-NaN fill, and its DataWarning, only when context == "fit".
  • X_val is passed with context="validation" in both _FitMixin._build_model and SklearnBaseLSS.build_model.
  • InferenceModel.validate_input() passes context="predict".

Validation and prediction inputs keep their missing values, and the preprocessor's fitted imputer fills them. Training data is unchanged: a training column that is entirely NaN is still filled with the constant. The imputer is then fitted on that constant, so later missing values in that column at prediction time are imputed with the same constant.

Tests

  • tests/core/test_exceptions.py: with context="predict" or "validation", all-NaN numeric and object columns are left as NaN and no DataWarning is emitted.
  • tests/integration/test_data_workflows.py: a single-row prediction with a missing numeric or categorical value equals the prediction for the same row inside a larger batch, and InferenceModel.validate_input() keeps the missing value.

All 4 new cases fail on main and pass with this change.

Docs

  • The ensure_dataframe docstring now says the fill is training-only.
  • The FAQ note on entirely-missing columns explains that validation and prediction inputs, including single-row requests, are imputed with the fitted strategy.
  • The README note now says "training column".

Checks

  • just test: 3240 passed, 46 skipped, 28 xfailed; coverage 91.29%
  • just docs: build succeeded (warnings treated as errors)
  • just check: all hooks passed (ruff, prettier, pyright)

ChrisW09 and others added 2 commits October 21, 2026 15:34
ensure_dataframe() filled every entirely-NaN column with 0 or "missing"
in every context. During prediction a single row with a missing value
forms an all-NaN column, so it bypassed the imputer fitted on the
training data and the prediction depended on the other rows in the
batch. The constant fill now applies only to training data
(context="fit"); validation input and prediction input, including
InferenceModel.validate_input(), keep their missing values for the
fitted preprocessor.

Fixes #495

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Missing values in single-row predictions are filled with 0/"missing" instead of the fitted imputer value

2 participants