address more comments

ashleyxuu · ashleyxuu · commit aa7b7b0c90e0 · 2024-05-21T22:41:44.000Z
diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py
@@ -28,7 +28,7 @@
 
 from bigframes import constants
 from bigframes.core import log_adapter
-from bigframes.ml import base, core, globals, preprocessing, utils
+from bigframes.ml import base, core, globals, impute, preprocessing, utils
 import bigframes.pandas as bpd
 
 _BQML_TRANSFROM_TYPE_MAPPING = types.MappingProxyType(
@@ -40,7 +40,7 @@
         "ML.BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.QUANTILE_BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.LABEL_ENCODER": preprocessing.LabelEncoder,
-        "ML.IMPUTER": preprocessing.SimpleImputer,
+        "ML.IMPUTER": impute.SimpleImputer,
     }
 )
 
@@ -59,7 +59,7 @@ def __init__(
         transformers: List[
             Tuple[
                 str,
-                preprocessing.PreprocessingType,
+                Union[preprocessing.PreprocessingType, impute.SimpleImputer],
                 Union[str, List[str]],
             ]
         ],
@@ -74,12 +74,14 @@ def __init__(
     @property
     def transformers_(
         self,
-    ) -> List[Tuple[str, preprocessing.PreprocessingType, str,]]:
+    ) -> List[
+        Tuple[str, Union[preprocessing.PreprocessingType, impute.SimpleImputer], str]
+    ]:
         """The collection of transformers as tuples of (name, transformer, column)."""
         result: List[
             Tuple[
                 str,
-                preprocessing.PreprocessingType,
+                Union[preprocessing.PreprocessingType, impute.SimpleImputer],
                 str,
             ]
         ] = []
@@ -108,7 +110,7 @@ def _extract_from_bq_model(
         transformers: List[
             Tuple[
                 str,
-                preprocessing.PreprocessingType,
+                Union[preprocessing.PreprocessingType, impute.SimpleImputer],
                 Union[str, List[str]],
             ]
         ] = []
@@ -153,7 +155,9 @@ def camel_to_snake(name):
 
     def _merge(
         self, bq_model: bigquery.Model
-    ) -> Union[ColumnTransformer, preprocessing.PreprocessingType,]:
+    ) -> Union[
+        ColumnTransformer, Union[preprocessing.PreprocessingType, impute.SimpleImputer]
+    ]:
         """Try to merge the column transformer to a simple transformer. Depends on all the columns in bq_model are transformed with the same transformer."""
         transformers = self.transformers_
 
diff --git a/bigframes/ml/impute.py b/bigframes/ml/impute.py
@@ -0,0 +1,123 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformers for missing value imputation. This module is styled after
+scikit-learn's preprocessing module: https://scikit-learn.org/stable/modules/preprocessing.html."""
+
+from __future__ import annotations
+
+import typing
+from typing import Any, List, Literal, Optional, Tuple, Union
+
+import bigframes_vendored.sklearn.impute._base
+
+from bigframes.core import log_adapter
+from bigframes.ml import base, core, globals, utils
+import bigframes.pandas as bpd
+
+
+@log_adapter.class_logger
+class SimpleImputer(
+    base.Transformer,
+    bigframes_vendored.sklearn.impute._base.SimpleImputer,
+):
+
+    __doc__ = bigframes_vendored.sklearn.impute._base.SimpleImputer.__doc__
+
+    def __init__(
+        self,
+        strategy: Literal["mean", "median", "most_frequent"] = "mean",
+    ):
+        self.strategy = strategy
+        self._bqml_model: Optional[core.BqmlModel] = None
+        self._bqml_model_factory = globals.bqml_model_factory()
+        self._base_sql_generator = globals.base_sql_generator()
+
+    # TODO(garrettwu): implement __hash__
+    def __eq__(self, other: Any) -> bool:
+        return (
+            type(other) is SimpleImputer
+            and self.strategy == other.strategy
+            and self._bqml_model == other._bqml_model
+        )
+
+    def _compile_to_sql(
+        self,
+        columns: List[str],
+        X=None,
+    ) -> List[Tuple[str, str]]:
+        """Compile this transformer to a list of SQL expressions that can be included in
+        a BQML TRANSFORM clause
+
+        Args:
+            columns:
+                a list of column names to transform
+            X:
+                The Dataframe with training data.
+
+        Returns: a list of tuples of (sql_expression, output_name)"""
+        return [
+            (
+                self._base_sql_generator.ml_imputer(
+                    column, self.strategy, f"imputer_{column}"
+                ),
+                f"imputer_{column}",
+            )
+            for column in columns
+        ]
+
+    @classmethod
+    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(SimpleImputer, column_label)"""
+        s = sql[sql.find("(") + 1 : sql.find(")")]
+        col_label, strategy = s.split(", ")
+        return cls(strategy[1:-1]), col_label  # type: ignore
+
+    def fit(
+        self,
+        X: Union[bpd.DataFrame, bpd.Series],
+        y=None,  # ignored
+    ) -> SimpleImputer:
+        (X,) = utils.convert_to_dataframe(X)
+
+        compiled_transforms = self._compile_to_sql(X.columns.tolist(), X)
+        transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
+
+        self._bqml_model = self._bqml_model_factory.create_model(
+            X,
+            options={"model_type": "transform_only"},
+            transforms=transform_sqls,
+        )
+
+        # The schema of TRANSFORM output is not available in the model API, so save it during fitting
+        self._output_names = [name for _, name in compiled_transforms]
+        return self
+
+    def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
+        if not self._bqml_model:
+            raise RuntimeError("Must be fitted before transform")
+
+        (X,) = utils.convert_to_dataframe(X)
+
+        df = self._bqml_model.transform(X)
+        return typing.cast(
+            bpd.DataFrame,
+            df[self._output_names],
+        )
diff --git a/bigframes/ml/loader.py b/bigframes/ml/loader.py
@@ -29,6 +29,7 @@
     ensemble,
     forecasting,
     imported,
+    impute,
     linear_model,
     llm,
     pipeline,
@@ -84,6 +85,7 @@ def from_bq(
     pipeline.Pipeline,
     compose.ColumnTransformer,
     preprocessing.PreprocessingType,
+    impute.SimpleImputer,
 ]:
     """Load a BQML model to BigQuery DataFrames ML.
 
diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py
@@ -26,7 +26,15 @@
 import bigframes
 import bigframes.constants as constants
 from bigframes.core import log_adapter
-from bigframes.ml import base, compose, forecasting, loader, preprocessing, utils
+from bigframes.ml import (
+    base,
+    compose,
+    forecasting,
+    impute,
+    loader,
+    preprocessing,
+    utils,
+)
 import bigframes.pandas as bpd
 
 
@@ -56,6 +64,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]):
                 preprocessing.MinMaxScaler,
                 preprocessing.KBinsDiscretizer,
                 preprocessing.LabelEncoder,
+                impute.SimpleImputer,
             ),
         ):
             self._transform = transform
diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py
@@ -20,7 +20,6 @@
 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
@@ -416,102 +415,6 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
         )
 
 
-@log_adapter.class_logger
-class SimpleImputer(
-    base.Transformer,
-    bigframes_vendored.sklearn.impute._base.SimpleImputer,
-):
-
-    __doc__ = bigframes_vendored.sklearn.impute._base.SimpleImputer.__doc__
-
-    def __init__(
-        self,
-        strategy: Literal["mean", "median", "most_frequent"] = "mean",
-    ):
-        self.strategy = strategy
-        self._bqml_model: Optional[core.BqmlModel] = None
-        self._bqml_model_factory = globals.bqml_model_factory()
-        self._base_sql_generator = globals.base_sql_generator()
-
-    # TODO(garrettwu): implement __hash__
-    def __eq__(self, other: Any) -> bool:
-        return (
-            type(other) is SimpleImputer
-            and self.strategy == other.strategy
-            and self._bqml_model == other._bqml_model
-        )
-
-    def _compile_to_sql(
-        self,
-        columns: List[str],
-        X=None,
-    ) -> List[Tuple[str, str]]:
-        """Compile this transformer to a list of SQL expressions that can be included in
-        a BQML TRANSFORM clause
-
-        Args:
-            columns:
-                a list of column names to transform
-            X:
-                The Dataframe with training data.
-
-        Returns: a list of tuples of (sql_expression, output_name)"""
-        return [
-            (
-                self._base_sql_generator.ml_imputer(
-                    column, self.strategy, f"imputer_{column}"
-                ),
-                f"imputer_{column}",
-            )
-            for column in columns
-        ]
-
-    @classmethod
-    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(SimpleImputer, column_label)"""
-        s = sql[sql.find("(") + 1 : sql.find(")")]
-        col_label, strategy = s.split(", ")
-        return cls(strategy[1:-1]), col_label  # type: ignore
-
-    def fit(
-        self,
-        X: Union[bpd.DataFrame, bpd.Series],
-        y=None,  # ignored
-    ) -> SimpleImputer:
-        (X,) = utils.convert_to_dataframe(X)
-
-        compiled_transforms = self._compile_to_sql(X.columns.tolist(), X)
-        transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
-
-        self._bqml_model = self._bqml_model_factory.create_model(
-            X,
-            options={"model_type": "transform_only"},
-            transforms=transform_sqls,
-        )
-
-        # The schema of TRANSFORM output is not available in the model API, so save it during fitting
-        self._output_names = [name for _, name in compiled_transforms]
-        return self
-
-    def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
-        if not self._bqml_model:
-            raise RuntimeError("Must be fitted before transform")
-
-        (X,) = utils.convert_to_dataframe(X)
-
-        df = self._bqml_model.transform(X)
-        return typing.cast(
-            bpd.DataFrame,
-            df[self._output_names],
-        )
-
-
 @log_adapter.class_logger
 class OneHotEncoder(
     base.Transformer,
@@ -765,5 +668,4 @@ def transform(self, y: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
     MinMaxScaler,
     KBinsDiscretizer,
     LabelEncoder,
-    SimpleImputer,
 ]
diff --git a/docs/reference/bigframes.ml/impute.rst b/docs/reference/bigframes.ml/impute.rst
@@ -0,0 +1,7 @@
+bigframes.ml.impute
+==========================
+
+.. automodule:: bigframes.ml.impute
+    :members:
+    :inherited-members:
+    :undoc-members:
diff --git a/docs/reference/bigframes.ml/index.rst b/docs/reference/bigframes.ml/index.rst
@@ -19,6 +19,8 @@ API Reference
 
     imported
 
+    impute
+
     linear_model
 
     llm
diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml
@@ -134,6 +134,12 @@
       - name: XGBoostModel
         uid: bigframes.ml.imported.XGBoostModel
       name: imported
+    - items:
+      - name: Overview
+        uid: bigframes.ml.impute
+      - name: SimpleImputer
+        uid: bigframes.ml.impute.SimpleImputer
+      name: impute
     - items:
       - name: Overview
         uid: bigframes.ml.linear_model
diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py
diff --git a/tests/system/small/ml/test_impute.py b/tests/system/small/ml/test_impute.py
diff --git a/tests/system/small/ml/test_preprocessing.py b/tests/system/small/ml/test_preprocessing.py
diff --git a/third_party/bigframes_vendored/sklearn/impute/_base.py b/third_party/bigframes_vendored/sklearn/impute/_base.py
