address comments

Author: ashleyxuu

bigframes/ml/compose.py

@@ -40,7 +40,7 @@
         "ML.BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.QUANTILE_BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.LABEL_ENCODER": preprocessing.LabelEncoder,
-        "ML.IMPUTER": preprocessing.Imputer,
+        "ML.IMPUTER": preprocessing.SimpleImputer,
     }
 )
 

bigframes/ml/preprocessing.py

@@ -20,10 +20,10 @@
 import typing
 from typing import Any, cast, List, Literal, Optional, Tuple, Union
 
+import bigframes_vendored.sklearn.impute._base
 import bigframes_vendored.sklearn.preprocessing._data
 import bigframes_vendored.sklearn.preprocessing._discretization
 import bigframes_vendored.sklearn.preprocessing._encoder
-import bigframes_vendored.sklearn.preprocessing._imputation
 import bigframes_vendored.sklearn.preprocessing._label
 
 from bigframes.core import log_adapter
@@ -417,12 +417,12 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
 
 
 @log_adapter.class_logger
-class Imputer(
+class SimpleImputer(
     base.Transformer,
-    bigframes_vendored.sklearn.preprocessing._imputation.Imputer,
+    bigframes_vendored.sklearn.impute._base.SimpleImputer,
 ):
 
-    __doc__ = bigframes_vendored.sklearn.preprocessing._imputation.Imputer.__doc__
+    __doc__ = bigframes_vendored.sklearn.impute._base.SimpleImputer.__doc__
 
     def __init__(
         self,
@@ -436,7 +436,7 @@ def __init__(
     # TODO(garrettwu): implement __hash__
     def __eq__(self, other: Any) -> bool:
         return (
-            type(other) is Imputer
+            type(other) is SimpleImputer
             and self.strategy == other.strategy
             and self._bqml_model == other._bqml_model
         )
@@ -467,14 +467,14 @@ def _compile_to_sql(
         ]
 
     @classmethod
-    def _parse_from_sql(cls, sql: str) -> tuple[Imputer, str]:
-        """Parse SQL to tuple(Imputer, column_label).
+    def _parse_from_sql(cls, sql: str) -> tuple[SimpleImputer, str]:
+        """Parse SQL to tuple(SimpleImputer, column_label).
 
         Args:
             sql: SQL string of format "ML.IMPUTER({col_label}, {strategy}) OVER()"
 
         Returns:
-            tuple(Imputer, column_label)"""
+            tuple(SimpleImputer, column_label)"""
         s = sql[sql.find("(") + 1 : sql.find(")")]
         col_label, strategy = s.split(", ")
         return cls(strategy[1:-1]), col_label  # type: ignore
@@ -483,7 +483,7 @@ def fit(
         self,
         X: Union[bpd.DataFrame, bpd.Series],
         y=None,  # ignored
-    ) -> Imputer:
+    ) -> SimpleImputer:
         (X,) = utils.convert_to_dataframe(X)
 
         compiled_transforms = self._compile_to_sql(X.columns.tolist(), X)
@@ -765,5 +765,5 @@ def transform(self, y: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
     MinMaxScaler,
     KBinsDiscretizer,
     LabelEncoder,
-    Imputer,
+    SimpleImputer,
 ]

tests/system/small/ml/test_preprocessing.py

@@ -548,15 +548,15 @@ def test_k_bins_discretizer_save_load_quantile(new_penguins_df, dataset_id):
     assert reloaded_transformer._bqml_model is not None
 
 
-def test_imputer_normalized_fit_transform_default_params():
+def test_simple_imputer_normalized_fit_transform_default_params():
     missing_df = bpd.DataFrame(
         {
             "culmen_length_mm": [39.5, 38.5, 37.9],
             "culmen_depth_mm": [np.nan, 17.2, 18.1],
             "flipper_length_mm": [np.nan, 181.0, 188.0],
         }
     )
-    imputer = preprocessing.Imputer(strategy="mean")
+    imputer = preprocessing.SimpleImputer(strategy="mean")
     result = imputer.fit_transform(
         missing_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]]
     ).to_pandas()
@@ -574,15 +574,15 @@ def test_imputer_normalized_fit_transform_default_params():
     pd.testing.assert_frame_equal(result, expected)
 
 
-def test_imputer_series_normalizes(new_penguins_df):
+def test_simple_imputer_series_normalizes(new_penguins_df):
     missing_df = bpd.DataFrame(
         {
             "culmen_length_mm": [39.5, 38.5, 37.9],
             "culmen_depth_mm": [np.nan, 17.2, 18.1],
             "flipper_length_mm": [np.nan, 181.0, 188.0],
         }
     )
-    imputer = preprocessing.Imputer()
+    imputer = preprocessing.SimpleImputer()
     imputer.fit(missing_df["culmen_depth_mm"])
 
     result = imputer.transform(missing_df["culmen_depth_mm"]).to_pandas()
@@ -599,44 +599,44 @@ def test_imputer_series_normalizes(new_penguins_df):
     pd.testing.assert_frame_equal(result, expected, rtol=0.1)
 
 
-def test_imputer_save_load_mean(dataset_id):
+def test_simple_imputer_save_load_mean(dataset_id):
     missing_df = bpd.DataFrame(
         {
             "culmen_length_mm": [39.5, 38.5, 37.9],
             "culmen_depth_mm": [np.nan, 17.2, 18.1],
             "flipper_length_mm": [np.nan, 181.0, 188.0],
         }
     )
-    transformer = preprocessing.Imputer(strategy="mean")
+    transformer = preprocessing.SimpleImputer(strategy="mean")
     transformer.fit(
         missing_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]]
     )
 
     reloaded_transformer = transformer.to_gbq(
         f"{dataset_id}.temp_configured_model", replace=True
     )
-    assert isinstance(reloaded_transformer, preprocessing.Imputer)
+    assert isinstance(reloaded_transformer, preprocessing.SimpleImputer)
     assert reloaded_transformer.strategy == transformer.strategy
     assert reloaded_transformer._bqml_model is not None
 
 
-def test_imputer_save_load_most_frequent(dataset_id):
+def test_simple_imputer_save_load_most_frequent(dataset_id):
     missing_df = bpd.DataFrame(
         {
             "culmen_length_mm": [39.5, 38.5, 37.9],
             "culmen_depth_mm": [np.nan, 17.2, 18.1],
             "flipper_length_mm": [np.nan, 181.0, 188.0],
         }
     )
-    transformer = preprocessing.Imputer(strategy="most_frequent")
+    transformer = preprocessing.SimpleImputer(strategy="most_frequent")
     transformer.fit(
         missing_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]]
     )
 
     reloaded_transformer = transformer.to_gbq(
         f"{dataset_id}.temp_configured_model", replace=True
     )
-    assert isinstance(reloaded_transformer, preprocessing.Imputer)
+    assert isinstance(reloaded_transformer, preprocessing.SimpleImputer)
     assert reloaded_transformer.strategy == transformer.strategy
     assert reloaded_transformer._bqml_model is not None
 

third_party/bigframes_vendored/sklearn/preprocessing/_imputation.py renamed to third_party/bigframes_vendored/sklearn/impute/_base.py

@@ -1,13 +1,25 @@
 # Authors: Nicolas Tresegnie <[email protected]>
+#          Sergey Feldman <[email protected]>
 # License: BSD 3 clause
 
 from bigframes_vendored.sklearn.base import BaseEstimator, TransformerMixin
 
 from bigframes import constants
 
 
-class Imputer(BaseEstimator, TransformerMixin):
-    """Imputation transformer for completing missing values.
+class _BaseImputer(TransformerMixin, BaseEstimator):
+    """Base class for all imputers.
+
+    It adds automatically support for `add_indicator`.
+    """
+
+
+class SimpleImputer(_BaseImputer):
+    """
+    Univariate imputer for completing missing values with simple strategies.
+
+    Replace missing values using a descriptive statistic (e.g. mean, median, or
+    most frequent) along each column, or using a constant value.
 
     Args:
         strategy ({'mean', 'median', 'most_frequent'}, default='mean'):