feat: add ml.preprocessing.KBinsDiscretizer (#81)

* feat: add ml.preprocessing.KBinsDiscretizer

* fix: address all the comments

* fix: address additional comments

* fix: fix the failed test

* Empty commit

* Trigger Kokoro

Author: ashleyxuu

bigframes/ml/compose.py

@@ -31,6 +31,7 @@
     preprocessing.StandardScaler,
     preprocessing.MaxAbsScaler,
     preprocessing.MinMaxScaler,
+    preprocessing.KBinsDiscretizer,
     preprocessing.LabelEncoder,
 ]
 
@@ -91,18 +92,24 @@ def transformers_(
 
         return result
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    def _compile_to_sql(
+        self,
+        columns: List[str],
+        X: bpd.DataFrame,
+    ) -> List[Tuple[str, str]]:
         """Compile this transformer to a list of SQL expressions that can be included in
         a BQML TRANSFORM clause
 
         Args:
             columns (List[str]):
                 a list of column names to transform
+            X (bpd.DataFrame):
+                The Dataframe with training data.
 
         Returns:
             a list of tuples of (sql_expression, output_name)"""
         return [
-            transformer._compile_to_sql([column])[0]
+            transformer._compile_to_sql([column], X=X)[0]
             for column in columns
             for _, transformer, target_column in self.transformers_
             if column == target_column
@@ -115,7 +122,7 @@ def fit(
     ) -> ColumnTransformer:
         (X,) = utils.convert_to_dataframe(X)
 
-        compiled_transforms = self._compile_to_sql(X.columns.tolist())
+        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(

bigframes/ml/pipeline.py

@@ -52,6 +52,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]):
                 preprocessing.OneHotEncoder,
                 preprocessing.MaxAbsScaler,
                 preprocessing.MinMaxScaler,
+                preprocessing.KBinsDiscretizer,
                 preprocessing.LabelEncoder,
             ),
         ):
@@ -93,7 +94,7 @@ def fit(
     ) -> Pipeline:
         (X,) = utils.convert_to_dataframe(X)
 
-        compiled_transforms = self._transform._compile_to_sql(X.columns.tolist())
+        compiled_transforms = self._transform._compile_to_sql(X.columns.tolist(), X=X)
         transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms]
 
         if y is not None:
@@ -151,6 +152,7 @@ def _extract_as_column_transformer(
                 preprocessing.StandardScaler,
                 preprocessing.MaxAbsScaler,
                 preprocessing.MinMaxScaler,
+                preprocessing.KBinsDiscretizer,
                 preprocessing.LabelEncoder,
             ],
             Union[str, List[str]],
@@ -190,6 +192,13 @@ def _extract_as_column_transformer(
                     *preprocessing.MinMaxScaler._parse_from_sql(transform_sql),
                 )
             )
+        elif transform_sql.startswith("ML.BUCKETIZE"):
+            transformers.append(
+                (
+                    "k_bins_discretizer",
+                    *preprocessing.KBinsDiscretizer._parse_from_sql(transform_sql),
+                )
+            )
         elif transform_sql.startswith("ML.LABEL_ENCODER"):
             transformers.append(
                 (
@@ -213,6 +222,7 @@ def _merge_column_transformer(
     preprocessing.OneHotEncoder,
     preprocessing.MaxAbsScaler,
     preprocessing.MinMaxScaler,
+    preprocessing.KBinsDiscretizer,
     preprocessing.LabelEncoder,
 ]:
     """Try to merge the column transformer to a simple transformer."""

bigframes/ml/preprocessing.py

@@ -23,6 +23,7 @@
 from bigframes.ml import base, core, globals, utils
 import bigframes.pandas as bpd
 import third_party.bigframes_vendored.sklearn.preprocessing._data
+import third_party.bigframes_vendored.sklearn.preprocessing._discretization
 import third_party.bigframes_vendored.sklearn.preprocessing._encoder
 import third_party.bigframes_vendored.sklearn.preprocessing._label
 
@@ -44,12 +45,15 @@ def __init__(self):
     def __eq__(self, other: Any) -> bool:
         return type(other) is StandardScaler and self._bqml_model == other._bqml_model
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    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
+            columns:
+                a list of column names to transform.
+            X (default None):
+                Ignored.
 
         Returns: a list of tuples of (sql_expression, output_name)"""
         return [
@@ -124,12 +128,15 @@ def __init__(self):
     def __eq__(self, other: Any) -> bool:
         return type(other) is MaxAbsScaler and self._bqml_model == other._bqml_model
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    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
+            columns:
+                a list of column names to transform.
+            X (default None):
+                Ignored.
 
         Returns: a list of tuples of (sql_expression, output_name)"""
         return [
@@ -204,12 +211,15 @@ def __init__(self):
     def __eq__(self, other: Any) -> bool:
         return type(other) is MinMaxScaler and self._bqml_model == other._bqml_model
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    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
+            columns:
+                a list of column names to transform.
+            X (default None):
+                Ignored.
 
         Returns: a list of tuples of (sql_expression, output_name)"""
         return [
@@ -267,6 +277,124 @@ def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
         )
 
 
+class KBinsDiscretizer(
+    base.Transformer,
+    third_party.bigframes_vendored.sklearn.preprocessing._discretization.KBinsDiscretizer,
+):
+    __doc__ = (
+        third_party.bigframes_vendored.sklearn.preprocessing._discretization.KBinsDiscretizer.__doc__
+    )
+
+    def __init__(
+        self,
+        n_bins: int = 5,
+        strategy: Literal["uniform", "quantile"] = "quantile",
+    ):
+        if strategy != "uniform":
+            raise NotImplementedError(
+                f"Only strategy = 'uniform' is supported now, input is {strategy}."
+            )
+        if n_bins < 2:
+            raise ValueError(
+                f"n_bins has to be larger than or equal to 2, input is {n_bins}."
+            )
+        self.n_bins = n_bins
+        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 KBinsDiscretizer
+            and self.n_bins == other.n_bins
+            and self._bqml_model == other._bqml_model
+        )
+
+    def _compile_to_sql(
+        self,
+        columns: List[str],
+        X: bpd.DataFrame,
+    ) -> 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)"""
+        array_split_points = {}
+        if self.strategy == "uniform":
+            for column in columns:
+                min_value = X[column].min()
+                max_value = X[column].max()
+                bin_size = (max_value - min_value) / self.n_bins
+                array_split_points[column] = [
+                    min_value + i * bin_size for i in range(self.n_bins - 1)
+                ]
+
+        return [
+            (
+                self._base_sql_generator.ml_bucketize(
+                    column, array_split_points[column], f"kbinsdiscretizer_{column}"
+                ),
+                f"kbinsdiscretizer_{column}",
+            )
+            for column in columns
+        ]
+
+    @classmethod
+    def _parse_from_sql(cls, sql: str) -> tuple[KBinsDiscretizer, str]:
+        """Parse SQL to tuple(KBinsDiscretizer, column_label).
+
+        Args:
+            sql: SQL string of format "ML.BUCKETIZE({col_label}, array_split_points, FALSE) OVER()"
+
+        Returns:
+            tuple(KBinsDiscretizer, column_label)"""
+        s = sql[sql.find("(") + 1 : sql.find(")")]
+        array_split_points = s[s.find("[") + 1 : s.find("]")]
+        col_label = s[: s.find(",")]
+        n_bins = array_split_points.count(",") + 2
+        return cls(n_bins, "uniform"), col_label
+
+    def fit(
+        self,
+        X: Union[bpd.DataFrame, bpd.Series],
+        y=None,  # ignored
+    ) -> KBinsDiscretizer:
+        (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],
+        )
+
+
 class OneHotEncoder(
     base.Transformer,
     third_party.bigframes_vendored.sklearn.preprocessing._encoder.OneHotEncoder,
@@ -308,13 +436,15 @@ def __eq__(self, other: Any) -> bool:
             and self.max_categories == other.max_categories
         )
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    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
+                a list of column names to transform.
+            X (default None):
+                Ignored.
 
         Returns: a list of tuples of (sql_expression, output_name)"""
 
@@ -432,13 +562,15 @@ def __eq__(self, other: Any) -> bool:
             and self.max_categories == other.max_categories
         )
 
-    def _compile_to_sql(self, columns: List[str]) -> List[Tuple[str, str]]:
+    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
+                a list of column names to transform.
+            X (default None):
+                Ignored.
 
         Returns: a list of tuples of (sql_expression, output_name)"""
 

bigframes/ml/sql.py

@@ -85,6 +85,15 @@ def ml_min_max_scaler(self, numeric_expr_sql: str, name: str) -> str:
         """Encode ML.MIN_MAX_SCALER for BQML"""
         return f"""ML.MIN_MAX_SCALER({numeric_expr_sql}) OVER() AS {name}"""
 
+    def ml_bucketize(
+        self,
+        numeric_expr_sql: str,
+        array_split_points: Iterable[Union[int, float]],
+        name: str,
+    ) -> str:
+        """Encode ML.MIN_MAX_SCALER for BQML"""
+        return f"""ML.BUCKETIZE({numeric_expr_sql}, {array_split_points}, FALSE) AS {name}"""
+
     def ml_one_hot_encoder(
         self,
         numeric_expr_sql: str,