add missing_only=True to all imputers to use in combination with variables=None

feature_engine/_docstrings/init_parameters/imputers.py

@@ -0,0 +1,4 @@
+_missing_only_docstring = """missing_only: bool, default=False
+        If `True` and `variables` is `None`, the transformer will select and impute
+        only the variables that show missing values during `fit()`.
+    """.rstrip()

feature_engine/imputation/arbitrary_number.py

@@ -28,6 +28,9 @@
 from feature_engine._docstrings.init_parameters.all_transformers import (
     _return_empty_docstring
 )
+from feature_engine._docstrings.init_parameters.imputers import (
+    _missing_only_docstring,
+)
 from feature_engine._docstrings.substitute import Substitution
 from feature_engine.dataframe_checks import check_X
 from feature_engine.imputation.base_imputer import BaseImputer
@@ -41,6 +44,7 @@
     imputer_dict_=_imputer_dict_docstring,
     variables_=_variables_attribute_docstring,
     return_empty=_return_empty_docstring,
+    missing_only=_missing_only_docstring,
     feature_names_in_=_feature_names_in_docstring,
     n_features_in_=_n_features_in_docstring,
     fit=_fit_not_learn_docstring,
@@ -76,6 +80,8 @@ class ArbitraryNumberImputer(BaseImputer):
         The dictionary of variables and the arbitrary numbers for their imputation. If
         specified, it overrides the above parameters.
 
+    {missing_only}
+
 
     Attributes
     ----------
@@ -126,13 +132,16 @@ def __init__(
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         return_empty: bool = False,
         imputer_dict: Optional[dict] = None,
+        missing_only: bool = False,
     ) -> None:
 
         if isinstance(arbitrary_number, int) or isinstance(arbitrary_number, float):
             self.arbitrary_number = arbitrary_number
         else:
             raise ValueError("arbitrary_number must be numeric of type int or float")
 
+        super().__init__(missing_only)
+
         _check_numerical_dict(imputer_dict)
 
         self.variables = _check_variables_input_value(variables)
@@ -168,6 +177,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         else:
             if self.variables is None:
                 self.variables_ = find_numerical_variables(X, self.return_empty)
+                if self.missing_only:
+                    self.variables_ = self._filter_variables_with_na(X, self.variables_)
             else:
                 self.variables_ = check_numerical_variables(X, self.variables)
             self.imputer_dict_ = {var: self.arbitrary_number for var in self.variables_}

feature_engine/imputation/base_imputer.py

@@ -12,6 +12,21 @@
 class BaseImputer(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin):
     """shared set-up checks and methods across imputers"""
 
+    def __init__(
+        self,
+        missing_only: bool = False,
+    ) -> None:
+        if not isinstance(missing_only, bool):
+            raise ValueError(
+                f"missing_only must be a boolean. Got {missing_only} instead."
+            )
+
+        self.missing_only = missing_only
+
+    def _filter_variables_with_na(self, X: pd.DataFrame, variables):
+        """Return variables that contain missing values."""
+        return [var for var in variables if X[var].isnull().any()]
+
     def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
         """
         Common checks before transforming data:

feature_engine/imputation/categorical.py

@@ -24,6 +24,9 @@
 from feature_engine._docstrings.init_parameters.all_transformers import (
     _return_empty_docstring
 )
+from feature_engine._docstrings.init_parameters.imputers import (
+    _missing_only_docstring,
+)
 from feature_engine._docstrings.substitute import Substitution
 from feature_engine.dataframe_checks import check_X
 from feature_engine.imputation.base_imputer import BaseImputer
@@ -40,6 +43,7 @@
     imputer_dict_=_imputer_dict_docstring,
     variables_=_variables_attribute_docstring,
     return_empty=_return_empty_docstring,
+    missing_only=_missing_only_docstring,
     feature_names_in_=_feature_names_in_docstring,
     n_features_in_=_n_features_in_docstring,
     transform=_transform_imputers_docstring,
@@ -97,6 +101,8 @@ class CategoricalImputer(BaseImputer):
         type object or categorical. If True, the imputer will select all variables or
         accept all variables entered by the user, including those cast as numeric.
 
+    {missing_only}
+
     Attributes
     ----------
     {imputer_dict_}
@@ -145,6 +151,7 @@ def __init__(
         return_empty: bool = False,
         return_object: bool = False,
         ignore_format: bool = False,
+        missing_only: bool = False,
     ) -> None:
         if imputation_method not in ["missing", "frequent"]:
             raise ValueError(
@@ -154,6 +161,8 @@ def __init__(
         if not isinstance(ignore_format, bool):
             raise ValueError("ignore_format takes only booleans True and False")
 
+        super().__init__(missing_only)
+
         self.imputation_method = imputation_method
         self.fill_value = fill_value
         self.variables = _check_variables_input_value(variables)
@@ -190,6 +199,9 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
             else:
                 self.variables_ = check_categorical_variables(X, self.variables)
 
+        if self.variables is None and self.missing_only:
+            self.variables_ = self._filter_variables_with_na(X, self.variables_)
+
         if self.imputation_method == "missing":
             self.imputer_dict_ = {var: self.fill_value for var in self.variables_}
 

feature_engine/imputation/end_tail.py

@@ -20,6 +20,9 @@
 from feature_engine._docstrings.init_parameters.all_transformers import (
     _variables_numerical_docstring, _return_empty_docstring
 )
+from feature_engine._docstrings.init_parameters.imputers import (
+    _missing_only_docstring,
+)
 from feature_engine._docstrings.methods import (
     _fit_transform_docstring,
     _transform_imputers_docstring,
@@ -36,6 +39,7 @@
 @Substitution(
     variables=_variables_numerical_docstring,
     return_empty=_return_empty_docstring,
+    missing_only=_missing_only_docstring,
     imputer_dict_=_imputer_dict_docstring,
     variables_=_variables_attribute_docstring,
     feature_names_in_=_feature_names_in_docstring,
@@ -105,6 +109,8 @@ class EndTailImputer(BaseImputer):
 
     {return_empty}
 
+    {missing_only}
+
     Attributes
     ----------
     {imputer_dict_}
@@ -149,6 +155,7 @@ def __init__(
         fold: int = 3,
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         return_empty: bool = False,
+        missing_only: bool = False,
     ) -> None:
 
         if imputation_method not in ["gaussian", "iqr", "max"]:
@@ -162,6 +169,8 @@ def __init__(
         if fold <= 0:
             raise ValueError("fold takes only positive numbers")
 
+        super().__init__(missing_only)
+
         self.imputation_method = imputation_method
         self.tail = tail
         self.fold = fold
@@ -188,6 +197,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # find or check for numerical variables
         if self.variables is None:
             self.variables_ = find_numerical_variables(X, self.return_empty)
+            if self.missing_only:
+                self.variables_ = self._filter_variables_with_na(X, self.variables_)
         else:
             self.variables_ = check_numerical_variables(X, self.variables)
 

feature_engine/imputation/mean_median.py

@@ -20,6 +20,9 @@
 from feature_engine._docstrings.init_parameters.all_transformers import (
     _variables_numerical_docstring, _return_empty_docstring
 )
+from feature_engine._docstrings.init_parameters.imputers import (
+    _missing_only_docstring,
+)
 from feature_engine._docstrings.methods import (
     _fit_transform_docstring,
     _transform_imputers_docstring,
@@ -36,6 +39,7 @@
 @Substitution(
     variables=_variables_numerical_docstring,
     return_empty=_return_empty_docstring,
+    missing_only=_missing_only_docstring,
     imputer_dict_=_imputer_dict_docstring,
     variables_=_variables_attribute_docstring,
     feature_names_in_=_feature_names_in_docstring,
@@ -63,6 +67,8 @@ class MeanMedianImputer(BaseImputer):
 
     {return_empty}
 
+    {missing_only}
+
     Attributes
     ----------
     {imputer_dict_}
@@ -108,11 +114,14 @@ def __init__(
         imputation_method: str = "median",
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         return_empty: bool = False,
+        missing_only: bool = False,
     ) -> None:
 
         if imputation_method not in ["median", "mean"]:
             raise ValueError("imputation_method takes only values 'median' or 'mean'")
 
+        super().__init__(missing_only)
+
         self.imputation_method = imputation_method
         self.variables = _check_variables_input_value(variables)
 
@@ -138,6 +147,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # find or check for numerical variables
         if self.variables is None:
             self.variables_ = find_numerical_variables(X, self.return_empty)
+            if self.missing_only:
+                self.variables_ = self._filter_variables_with_na(X, self.variables_)
         else:
             self.variables_ = check_numerical_variables(X, self.variables)
 

feature_engine/imputation/random_sample.py

@@ -24,6 +24,9 @@
 from feature_engine._docstrings.init_parameters.all_transformers import (
     _return_empty_docstring
 )
+from feature_engine._docstrings.init_parameters.imputers import (
+    _missing_only_docstring,
+)
 from feature_engine._docstrings.substitute import Substitution
 from feature_engine.dataframe_checks import check_X
 from feature_engine.imputation.base_imputer import BaseImputer
@@ -50,6 +53,7 @@ def _define_seed(
 @Substitution(
     variables_=_variables_attribute_docstring,
     return_empty=_return_empty_docstring,
+    missing_only=_missing_only_docstring,
     feature_names_in_=_feature_names_in_docstring,
     n_features_in_=_n_features_in_docstring,
     transform=_transform_imputers_docstring,
@@ -100,6 +104,8 @@ class RandomSampleImputer(BaseImputer):
         observation, you can choose to combine those values as an addition or a
         multiplication. Can take the values 'add' or 'multiply'.
 
+    {missing_only}
+
     Attributes
     ----------
     X_:
@@ -148,6 +154,7 @@ def __init__(
         random_state: Union[None, int, str, List[Union[str, int]]] = None,
         seed: str = "general",
         seeding_method: str = "add",
+        missing_only: bool = False,
     ) -> None:
 
         if seed not in ["general", "observation"]:
@@ -168,6 +175,8 @@ def __init__(
                 "or more variables which will be used to seed the imputer"
             )
 
+        super().__init__(missing_only)
+
         self.variables = _check_variables_input_value(variables)
 
         _check_return_empty_is_bool(return_empty)
@@ -199,6 +208,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # find variables to impute
         if self.variables is None:
             self.variables_ = find_all_variables(X, self.return_empty)
+            if self.missing_only:
+                self.variables_ = self._filter_variables_with_na(X, self.variables_)
         else:
             self.variables_ = check_all_variables(X, self.variables)
 

tests/test_imputation/test_arbitrary_number_imputer.py

@@ -82,3 +82,32 @@ def test_dictionary_of_imputation_values(df_na):
 def imputer_error_when_dictionary_value_is_string():
     with pytest.raises(ValueError):
         ArbitraryNumberImputer(imputer_dict={"Age": "arbitrary_number"})
+
+
+def test_missing_only_selects_numerical_variables_with_na(df_na):
+    imputer = ArbitraryNumberImputer(arbitrary_number=99, missing_only=True)
+    X = df_na.copy()
+    X["Var_No_Nulls"] = [333] * X.shape[0]
+    X_transformed = imputer.fit_transform(X)
+
+    expected_results_df = df_na.copy()
+    expected_results_df["Age"] = expected_results_df["Age"].fillna(99)
+    expected_results_df["Marks"] = expected_results_df["Marks"].fillna(99)
+    expected_results_df["Var_No_Nulls"] = [333] * X.shape[0]
+
+    assert imputer.variables_ == ["Age", "Marks"]
+    assert imputer.imputer_dict_ == {"Age": 99, "Marks": 99}
+    pd.testing.assert_frame_equal(X_transformed, expected_results_df)
+
+
+def test_missing_only_ignored_when_imputer_dict_provided(df_na):
+    imputer = ArbitraryNumberImputer(
+        imputer_dict={"Age": -42, "Marks": -999},
+        missing_only=True,
+    )
+    X = df_na.copy()
+    X["Marks"] = [0.9] * X.shape[0]
+    imputer.fit(X)
+
+    assert imputer.variables_ == ["Age", "Marks"]
+    assert imputer.imputer_dict_ == {"Age": -42, "Marks": -999}

tests/test_imputation/test_categorical_imputer.py

@@ -305,3 +305,39 @@ def test_error_when_ignore_format_is_not_boolean(ignore_format):
 
     # check that error message matches
     assert str(record.value) == msg
+
+
+def test_missing_only_selects_categorical_variables_with_na(df_na):
+    imputer = CategoricalImputer(
+        imputation_method="missing",
+        missing_only=True,
+    )
+    X = df_na.copy()
+    X["Cat_No_Nulls"] = ["pasta"] * X.shape[0]
+    X_transformed = imputer.fit_transform(X)
+
+    expected_results_df = df_na.copy()
+    expected_results_df["Name"] = expected_results_df["Name"].fillna("Missing")
+    expected_results_df["City"] = expected_results_df["City"].fillna("Missing")
+    expected_results_df["Studies"] = expected_results_df["Studies"].fillna("Missing")
+    expected_results_df["Cat_No_Nulls"] = ["pasta"] * X.shape[0]
+
+    assert imputer.variables_ == ["Name", "City", "Studies"]
+    assert "Cat_No_Nulls" not in imputer.imputer_dict_
+    pd.testing.assert_frame_equal(X_transformed, expected_results_df)
+
+
+def test_missing_only_with_ignore_format_selects_all_variables_with_na(df_na):
+    imputer = CategoricalImputer(
+        imputation_method="missing",
+        fill_value="Missing",
+        ignore_format=True,
+        missing_only=True,
+    )
+    X = df_na.copy()
+    X["Var_No_Nulls"] = [333] * X.shape[0]
+    imputer.fit(X)
+
+    assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks"]
+    assert "Var_No_Nulls" not in imputer.variables_
+    assert "dob" not in imputer.variables_

tests/test_imputation/test_end_tail_imputer.py

@@ -101,3 +101,24 @@ def test_error_when_tail_is_string():
 def test_error_when_fold_is_1():
     with pytest.raises(ValueError):
         EndTailImputer(fold=-1)
+
+
+def test_missing_only_selects_numerical_variables_with_na(df_na):
+    imputer = EndTailImputer(
+        imputation_method="iqr",
+        tail="right",
+        fold=1.5,
+        missing_only=True,
+    )
+    X = df_na.copy()
+    X["Var_No_Nulls"] = [1984] * X.shape[0]
+    X_transformed = imputer.fit_transform(X)
+
+    expected_results_df = df_na.copy()
+    expected_results_df["Age"] = expected_results_df["Age"].fillna(65.5)
+    expected_results_df["Marks"] = expected_results_df["Marks"].fillna(1.0625)
+    expected_results_df["Var_No_Nulls"] = [1984] * X.shape[0]
+
+    assert imputer.variables_ == ["Age", "Marks"]
+    assert "Var_No_Nulls" not in imputer.imputer_dict_
+    pd.testing.assert_frame_equal(X_transformed, expected_results_df)