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,
 ]