feat: add PolynomailFeatures to_gbq and pipeline support (#805)

* feat: add ml.preprocessing.PolynomialFeatures class

* feat: add PolynomailFeatures to_gbq and pipeline support

* fix tests

* fix tests

Author: GarrettWu

bigframes/ml/base.py

@@ -184,6 +184,16 @@ class BaseTransformer(BaseEstimator):
     def __init__(self):
         self._bqml_model: Optional[core.BqmlModel] = None
 
+    @abc.abstractmethod
+    def _keys(self):
+        pass
+
+    def __eq__(self, other) -> bool:
+        return type(self) is type(other) and self._keys() == other._keys()
+
+    def __hash__(self) -> int:
+        return hash(self._keys())
+
     _T = TypeVar("_T", bound="BaseTransformer")
 
     def to_gbq(self: _T, model_name: str, replace: bool = False) -> _T:

bigframes/ml/compose.py

@@ -21,7 +21,7 @@
 import re
 import types
 import typing
-from typing import cast, List, Optional, Tuple, Union
+from typing import cast, Iterable, List, Optional, Set, Tuple, Union
 
 import bigframes_vendored.sklearn.compose._column_transformer
 from google.cloud import bigquery
@@ -40,6 +40,7 @@
         "ML.BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.QUANTILE_BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.LABEL_ENCODER": preprocessing.LabelEncoder,
+        "ML.POLYNOMIAL_EXPAND": preprocessing.PolynomialFeatures,
         "ML.IMPUTER": impute.SimpleImputer,
     }
 )
@@ -56,21 +57,24 @@ class ColumnTransformer(
 
     def __init__(
         self,
-        transformers: List[
+        transformers: Iterable[
             Tuple[
                 str,
                 Union[preprocessing.PreprocessingType, impute.SimpleImputer],
-                Union[str, List[str]],
+                Union[str, Iterable[str]],
             ]
         ],
     ):
         # TODO: if any(transformers) has fitted raise warning
-        self.transformers = transformers
+        self.transformers = list(transformers)
         self._bqml_model: Optional[core.BqmlModel] = None
         self._bqml_model_factory = globals.bqml_model_factory()
         # call self.transformers_ to check chained transformers
         self.transformers_
 
+    def _keys(self):
+        return (self.transformers, self._bqml_model)
+
     @property
     def transformers_(
         self,
@@ -107,13 +111,13 @@ def _extract_from_bq_model(
         """Extract transformers as ColumnTransformer obj from a BQ Model. Keep the _bqml_model field as None."""
         assert "transformColumns" in bq_model._properties
 
-        transformers: List[
+        transformers_set: Set[
             Tuple[
                 str,
                 Union[preprocessing.PreprocessingType, impute.SimpleImputer],
                 Union[str, List[str]],
             ]
-        ] = []
+        ] = set()
 
         def camel_to_snake(name):
             name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
@@ -134,7 +138,7 @@ def camel_to_snake(name):
             for prefix in _BQML_TRANSFROM_TYPE_MAPPING:
                 if transform_sql.startswith(prefix):
                     transformer_cls = _BQML_TRANSFROM_TYPE_MAPPING[prefix]
-                    transformers.append(
+                    transformers_set.add(
                         (
                             camel_to_snake(transformer_cls.__name__),
                             *transformer_cls._parse_from_sql(transform_sql),  # type: ignore
@@ -148,7 +152,7 @@ def camel_to_snake(name):
                     f"Unsupported transformer type. {constants.FEEDBACK_LINK}"
                 )
 
-        transformer = cls(transformers=transformers)
+        transformer = cls(transformers=list(transformers_set))
         transformer._output_names = output_names
 
         return transformer
@@ -159,23 +163,37 @@ def _merge(
         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_
+        transformers = self.transformers
 
         assert len(transformers) > 0
         _, transformer_0, column_0 = transformers[0]
+        feature_columns_sorted = sorted(
+            [
+                cast(str, feature_column.name)
+                for feature_column in bq_model.feature_columns
+            ]
+        )
+
+        if (
+            len(transformers) == 1
+            and isinstance(transformer_0, preprocessing.PolynomialFeatures)
+            and sorted(column_0) == feature_columns_sorted
+        ):
+            transformer_0._output_names = self._output_names
+            return transformer_0
+
+        if not isinstance(column_0, str):
+            return self
         columns = [column_0]
         for _, transformer, column in transformers[1:]:
+            if not isinstance(column, str):
+                return self
             # all transformers are the same
             if transformer != transformer_0:
                 return self
             columns.append(column)
         # all feature columns are transformed
-        if sorted(
-            [
-                cast(str, feature_column.name)
-                for feature_column in bq_model.feature_columns
-            ]
-        ) == sorted(columns):
+        if sorted(columns) == feature_columns_sorted:
             transformer_0._output_names = self._output_names
             return transformer_0
 
@@ -197,12 +215,12 @@ def _compile_to_sql(
 
         Returns:
             a list of tuples of (sql_expression, output_name)"""
-        return [
-            transformer._compile_to_sql([column], X=X)[0]
-            for column in columns
-            for _, transformer, target_column in self.transformers_
-            if column == target_column
-        ]
+        result = []
+        for _, transformer, target_columns in self.transformers:
+            if isinstance(target_columns, str):
+                target_columns = [target_columns]
+            result += transformer._compile_to_sql(target_columns, X=X)
+        return result
 
     def fit(
         self,

bigframes/ml/core.py

@@ -112,6 +112,15 @@ def __init__(self, session: bigframes.Session, model: bigquery.Model):
             self.model_name
         )
 
+    def _keys(self):
+        return (self._session, self._model)
+
+    def __eq__(self, other):
+        return isinstance(other, self.__class__) and self._keys() == other._keys()
+
+    def __hash__(self):
+        return hash(self._keys())
+
     @property
     def session(self) -> bigframes.Session:
         """Get the BigQuery DataFrames session that this BQML model wrapper is tied to"""

bigframes/ml/impute.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Any, List, Literal, Optional, Tuple, Union
+from typing import Iterable, List, Literal, Optional, Tuple, Union
 
 import bigframes_vendored.sklearn.impute._base
 
@@ -44,17 +44,12 @@ def __init__(
         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 _keys(self):
+        return (self._bqml_model, self.strategy)
 
     def _compile_to_sql(
         self,
-        columns: List[str],
+        columns: Iterable[str],
         X=None,
     ) -> List[Tuple[str, str]]:
         """Compile this transformer to a list of SQL expressions that can be included in

bigframes/ml/pipeline.py

@@ -64,6 +64,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]):
                 preprocessing.MinMaxScaler,
                 preprocessing.KBinsDiscretizer,
                 preprocessing.LabelEncoder,
+                preprocessing.PolynomialFeatures,
                 impute.SimpleImputer,
             ),
         ):

bigframes/ml/preprocessing.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Any, cast, List, Literal, Optional, Tuple, Union
+from typing import cast, Iterable, List, Literal, Optional, Tuple, Union
 
 import bigframes_vendored.sklearn.preprocessing._data
 import bigframes_vendored.sklearn.preprocessing._discretization
@@ -43,11 +43,10 @@ def __init__(self):
         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 StandardScaler and self._bqml_model == other._bqml_model
+    def _keys(self):
+        return (self._bqml_model,)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -125,11 +124,10 @@ def __init__(self):
         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 MaxAbsScaler and self._bqml_model == other._bqml_model
+    def _keys(self):
+        return (self._bqml_model,)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -207,11 +205,10 @@ def __init__(self):
         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 MinMaxScaler and self._bqml_model == other._bqml_model
+    def _keys(self):
+        return (self._bqml_model,)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -301,18 +298,12 @@ def __init__(
         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 KBinsDiscretizer
-            and self.n_bins == other.n_bins
-            and self.strategy == other.strategy
-            and self._bqml_model == other._bqml_model
-        )
+    def _keys(self):
+        return (self._bqml_model, self.n_bins, self.strategy)
 
     def _compile_to_sql(
         self,
-        columns: List[str],
+        columns: Iterable[str],
         X: bpd.DataFrame,
     ) -> List[Tuple[str, str]]:
         """Compile this transformer to a list of SQL expressions that can be included in
@@ -446,17 +437,10 @@ def __init__(
         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 OneHotEncoder
-            and self._bqml_model == other._bqml_model
-            and self.drop == other.drop
-            and self.min_frequency == other.min_frequency
-            and self.max_categories == other.max_categories
-        )
+    def _keys(self):
+        return (self._bqml_model, self.drop, self.min_frequency, self.max_categories)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -572,16 +556,10 @@ def __init__(
         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 LabelEncoder
-            and self._bqml_model == other._bqml_model
-            and self.min_frequency == other.min_frequency
-            and self.max_categories == other.max_categories
-        )
+    def _keys(self):
+        return (self._bqml_model, self.min_frequency, self.max_categories)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -672,18 +650,17 @@ class PolynomialFeatures(
     )
 
     def __init__(self, degree: int = 2):
+        if degree not in range(1, 5):
+            raise ValueError(f"degree has to be [1, 4], input is {degree}.")
         self.degree = degree
         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 PolynomialFeatures and self._bqml_model == other._bqml_model
-        )
+    def _keys(self):
+        return (self._bqml_model, self.degree)
 
-    def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
+    def _compile_to_sql(self, columns: Iterable[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
 
@@ -705,17 +682,18 @@ def _compile_to_sql(self, columns: List[str], X=None) -> List[Tuple[str, str]]:
         ]
 
     @classmethod
-    def _parse_from_sql(cls, sql: str) -> tuple[PolynomialFeatures, str]:
-        """Parse SQL to tuple(PolynomialFeatures, column_label).
+    def _parse_from_sql(cls, sql: str) -> tuple[PolynomialFeatures, tuple[str, ...]]:
+        """Parse SQL to tuple(PolynomialFeatures, column_labels).
 
         Args:
             sql: SQL string of format "ML.POLYNOMIAL_EXPAND(STRUCT(col_label0, col_label1, ...), degree)"
 
         Returns:
             tuple(MaxAbsScaler, column_label)"""
-        col_label = sql[sql.find("STRUCT(") + 7 : sql.find(")")]
+        col_labels = sql[sql.find("STRUCT(") + 7 : sql.find(")")].split(",")
+        col_labels = [label.strip() for label in col_labels]
         degree = int(sql[sql.rfind(",") + 1 : sql.rfind(")")])
-        return cls(degree), col_label
+        return cls(degree), tuple(col_labels)
 
     def fit(
         self,
@@ -762,8 +740,6 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
             df[self._output_names],
         )
 
-    # TODO(garrettwu): to_gbq()
-
 
 PreprocessingType = Union[
     OneHotEncoder,
@@ -772,4 +748,5 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
     MinMaxScaler,
     KBinsDiscretizer,
     LabelEncoder,
+    PolynomialFeatures,
 ]