Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

feat: add PolynomialFeatures to_gbq and pipeline support #805

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jun 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions 10 bigframes/ml/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
60 changes: 39 additions & 21 deletions 60 bigframes/ml/compose.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
}
)
Expand All @@ -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,
Expand Down Expand Up @@ -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()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what's the reason for going from list to set?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The processor is different from others. Where previous are 1 -> 1 transform, poly_features is n -> m. So other transformers we get 1 transformed_column, but poly_features we get m transformed_column. Which need to be deduped. And the order of the transformers doesn't matter.


def camel_to_snake(name):
name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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

Expand All @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions 9 bigframes/ml/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down
13 changes: 4 additions & 9 deletions 13 bigframes/ml/impute.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions 1 bigframes/ml/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]):
preprocessing.MinMaxScaler,
preprocessing.KBinsDiscretizer,
preprocessing.LabelEncoder,
preprocessing.PolynomialFeatures,
impute.SimpleImputer,
),
):
Expand Down
83 changes: 30 additions & 53 deletions 83 bigframes/ml/preprocessing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -772,4 +748,5 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
MinMaxScaler,
KBinsDiscretizer,
LabelEncoder,
PolynomialFeatures,
]
Loading
Morty Proxy This is a proxified and sanitized view of the page, visit original site.