feat: add strategy="quantile" in KBinsDiscretizer (#654)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes internal #310685445 🦕

Author: ashleyxuu

bigframes/ml/compose.py

@@ -38,6 +38,7 @@
         "ML.MAX_ABS_SCALER": preprocessing.MaxAbsScaler,
         "ML.MIN_MAX_SCALER": preprocessing.MinMaxScaler,
         "ML.BUCKETIZE": preprocessing.KBinsDiscretizer,
+        "ML.QUANTILE_BUCKETIZE": preprocessing.KBinsDiscretizer,
         "ML.LABEL_ENCODER": preprocessing.LabelEncoder,
     }
 )

bigframes/ml/preprocessing.py

@@ -290,10 +290,6 @@ def __init__(
         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}."
@@ -337,30 +333,53 @@ def _compile_to_sql(
                     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}",
+            return [
+                (
+                    self._base_sql_generator.ml_bucketize(
+                        column, array_split_points[column], f"kbinsdiscretizer_{column}"
+                    ),
+                    f"kbinsdiscretizer_{column}",
+                )
+                for column in columns
+            ]
+
+        elif self.strategy == "quantile":
+
+            return [
+                (
+                    self._base_sql_generator.ml_quantile_bucketize(
+                        column, self.n_bins, f"kbinsdiscretizer_{column}"
+                    ),
+                    f"kbinsdiscretizer_{column}",
+                )
+                for column in columns
+            ]
+
+        else:
+            raise ValueError(
+                f"strategy should be set 'quantile' or 'uniform', but your input is {self.strategy}."
             )
-            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()"
+            sql: SQL string of format "ML.BUCKETIZE({col_label}, array_split_points, FALSE)"
+                or ML.QUANTILE_BUCKETIZE({col_label}, num_bucket) 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
+
+        if sql.startswith("ML.QUANTILE_BUCKETIZE"):
+            num_bins = s.split(",")[1]
+            return cls(int(num_bins), "quantile"), col_label
+        else:
+            array_split_points = s[s.find("[") + 1 : s.find("]")]
+            n_bins = array_split_points.count(",") + 2
+            return cls(n_bins, "uniform"), col_label
 
     def fit(
         self,

bigframes/ml/sql.py

@@ -109,9 +109,18 @@ def ml_bucketize(
         array_split_points: Iterable[Union[int, float]],
         name: str,
     ) -> str:
-        """Encode ML.MIN_MAX_SCALER for BQML"""
+        """Encode ML.BUCKETIZE for BQML"""
         return f"""ML.BUCKETIZE({numeric_expr_sql}, {array_split_points}, FALSE) AS {name}"""
 
+    def ml_quantile_bucketize(
+        self,
+        numeric_expr_sql: str,
+        num_bucket: int,
+        name: str,
+    ) -> str:
+        """Encode ML.QUANTILE_BUCKETIZE for BQML"""
+        return f"""ML.QUANTILE_BUCKETIZE({numeric_expr_sql}, {num_bucket}) OVER() AS {name}"""
+
     def ml_one_hot_encoder(
         self,
         numeric_expr_sql: str,

tests/system/small/ml/test_preprocessing.py

@@ -373,6 +373,27 @@ def test_k_bins_discretizer_normalized_fit_transform_default_params(new_penguins
     pd.testing.assert_frame_equal(result, expected, rtol=0.1)
 
 
+def test_k_bins_discretizer_normalized_fit_transform_default_params_quantile(
+    new_penguins_df,
+):
+    discretizer = preprocessing.KBinsDiscretizer(strategy="quantile")
+    result = discretizer.fit_transform(
+        new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]]
+    ).to_pandas()
+
+    expected = pd.DataFrame(
+        {
+            "kbinsdiscretizer_culmen_length_mm": ["bin_2", "bin_2", "bin_1"],
+            "kbinsdiscretizer_culmen_depth_mm": ["bin_2", "bin_1", "bin_2"],
+            "kbinsdiscretizer_flipper_length_mm": ["bin_2", "bin_1", "bin_2"],
+        },
+        dtype="string[pyarrow]",
+        index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
+    )
+
+    pd.testing.assert_frame_equal(result, expected, rtol=0.1)
+
+
 def test_k_bins_discretizer_series_normalizes(
     penguins_df_default_index, new_penguins_df
 ):
@@ -395,6 +416,28 @@ def test_k_bins_discretizer_series_normalizes(
     pd.testing.assert_frame_equal(result, expected, rtol=0.1)
 
 
+def test_k_bins_discretizer_series_normalizes_quantile(
+    penguins_df_default_index, new_penguins_df
+):
+    discretizer = preprocessing.KBinsDiscretizer(strategy="quantile")
+    discretizer.fit(penguins_df_default_index["culmen_length_mm"])
+
+    result = discretizer.transform(
+        penguins_df_default_index["culmen_length_mm"]
+    ).to_pandas()
+    result = discretizer.transform(new_penguins_df).to_pandas()
+
+    expected = pd.DataFrame(
+        {
+            "kbinsdiscretizer_culmen_length_mm": ["bin_2", "bin_2", "bin_1"],
+        },
+        dtype="string[pyarrow]",
+        index=pd.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
+    )
+
+    pd.testing.assert_frame_equal(result, expected, rtol=0.1)
+
+
 def test_k_bins_discretizer_normalizes(penguins_df_default_index, new_penguins_df):
     # TODO(http://b/292431644): add a second test that compares output to sklearn.preprocessing.KBinsDiscretizer, when BQML's change is in prod.
     discretizer = preprocessing.KBinsDiscretizer(strategy="uniform")
@@ -488,6 +531,21 @@ def test_k_bins_discretizer_save_load(new_penguins_df, dataset_id):
     pd.testing.assert_frame_equal(result, expected, rtol=0.1)
 
 
+def test_k_bins_discretizer_save_load_quantile(new_penguins_df, dataset_id):
+    transformer = preprocessing.KBinsDiscretizer(n_bins=6, strategy="quantile")
+    transformer.fit(
+        new_penguins_df[["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"]]
+    )
+
+    reloaded_transformer = transformer.to_gbq(
+        f"{dataset_id}.temp_configured_model", replace=True
+    )
+    assert isinstance(reloaded_transformer, preprocessing.KBinsDiscretizer)
+    assert reloaded_transformer.n_bins == transformer.n_bins
+    assert reloaded_transformer.strategy == transformer.strategy
+    assert reloaded_transformer._bqml_model is not None
+
+
 def test_one_hot_encoder_default_params(new_penguins_df):
     encoder = preprocessing.OneHotEncoder()
     encoder.fit(new_penguins_df[["species", "sex"]])

tests/unit/ml/test_sql.py

@@ -113,6 +113,13 @@ def test_k_bins_discretizer_correct(
     assert sql == "ML.BUCKETIZE(col_a, [1, 2, 3, 4], FALSE) AS scaled_col_a"
 
 
+def test_k_bins_discretizer_quantile_correct(
+    base_sql_generator: ml_sql.BaseSqlGenerator,
+):
+    sql = base_sql_generator.ml_quantile_bucketize("col_a", 5, "scaled_col_a")
+    assert sql == "ML.QUANTILE_BUCKETIZE(col_a, 5) OVER() AS scaled_col_a"
+
+
 def test_one_hot_encoder_correct(
     base_sql_generator: ml_sql.BaseSqlGenerator,
 ):

third_party/bigframes_vendored/sklearn/preprocessing/_discretization.py

@@ -18,7 +18,7 @@ class KBinsDiscretizer(TransformerMixin, BaseEstimator):
         strategy ({'uniform', 'quantile'}, default='quantile'):
             Strategy used to define the widths of the bins. 'uniform': All bins
             in each feature have identical widths. 'quantile': All bins in each
-            feature have the same number of points. Only `uniform` is supported.
+            feature have the same number of points.
     """
 
     def fit(self, X, y=None):