feat: add ml KMeans model params (#477)

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 #<issue_number_goes_here> 🦕

Author: GarrettWu

bigframes/ml/cluster.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import Dict, List, Optional, Union
+from typing import List, Literal, Optional, Union
 
 import bigframes_vendored.sklearn.cluster._kmeans
 from google.cloud import bigquery
@@ -27,6 +27,16 @@
 from bigframes.ml import base, core, globals, utils
 import bigframes.pandas as bpd
 
+_BQML_PARAMS_MAPPING = {
+    "n_clusters": "numClusters",
+    "init": "kmeansInitializationMethod",
+    "init_col": "kmeansInitializationColumn",
+    "distance_type": "distanceType",
+    "max_iter": "maxIterations",
+    "early_stop": "earlyStop",
+    "tol": "minRelativeProgress",
+}
+
 
 @log_adapter.class_logger
 class KMeans(
@@ -36,30 +46,67 @@ class KMeans(
 
     __doc__ = bigframes_vendored.sklearn.cluster._kmeans.KMeans.__doc__
 
-    def __init__(self, n_clusters: int = 8):
+    def __init__(
+        self,
+        n_clusters: int = 8,
+        *,
+        init: Literal["kmeans++", "random", "custom"] = "kmeans++",
+        init_col: Optional[str] = None,
+        distance_type: Literal["euclidean", "cosine"] = "euclidean",
+        max_iter: int = 20,
+        tol: float = 0.01,
+        warm_start: bool = False,
+    ):
         self.n_clusters = n_clusters
+        self.init = init
+        self.init_col = init_col
+        self.distance_type = distance_type
+        self.max_iter = max_iter
+        self.tol = tol
+        self.warm_start = warm_start
         self._bqml_model: Optional[core.BqmlModel] = None
         self._bqml_model_factory = globals.bqml_model_factory()
 
     @classmethod
     def _from_bq(cls, session: bigframes.Session, model: bigquery.Model) -> KMeans:
         assert model.model_type == "KMEANS"
 
-        kwargs = {}
+        kwargs: dict = {}
 
         # See https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun
         last_fitting = model.training_runs[-1]["trainingOptions"]
-        if "numClusters" in last_fitting:
-            kwargs["n_clusters"] = int(last_fitting["numClusters"])
+        dummy_kmeans = cls()
+        for bf_param, bf_value in dummy_kmeans.__dict__.items():
+            bqml_param = _BQML_PARAMS_MAPPING.get(bf_param)
+            if bqml_param in last_fitting:
+                # Convert types
+                kwargs[bf_param] = (
+                    str(last_fitting[bqml_param])
+                    if bf_param in ["init"]
+                    else type(bf_value)(last_fitting[bqml_param])
+                )
 
         new_kmeans = cls(**kwargs)
         new_kmeans._bqml_model = core.BqmlModel(session, model)
         return new_kmeans
 
     @property
-    def _bqml_options(self) -> Dict[str, str | int | float | List[str]]:
+    def _bqml_options(self) -> dict:
         """The model options as they will be set for BQML"""
-        return {"model_type": "KMEANS", "num_clusters": self.n_clusters}
+        options = {
+            "model_type": "KMEANS",
+            "num_clusters": self.n_clusters,
+            "KMEANS_INIT_METHOD": self.init,
+            "DISTANCE_TYPE": self.distance_type,
+            "MAX_ITERATIONS": self.max_iter,
+            "MIN_REL_PROGRESS": self.tol,
+            "WARM_START": self.warm_start,
+        }
+
+        if self.init_col is not None:
+            options["KMEANS_INIT_COL"] = self.init_col
+
+        return options
 
     def _fit(
         self,

tests/system/large/ml/test_cluster.py

@@ -19,11 +19,11 @@
 from tests.system.utils import assert_pandas_df_equal
 
 
-@pytest.mark.flaky(retries=2, delay=120)
+@pytest.mark.flaky(retries=2)
 def test_cluster_configure_fit_score_predict(
     session, penguins_df_default_index, dataset_id
 ):
-    model = cluster.KMeans(n_clusters=3)
+    model = cluster.KMeans(n_clusters=3, init="random")
 
     df = penguins_df_default_index.dropna()[
         [
@@ -118,3 +118,47 @@ def test_cluster_configure_fit_score_predict(
         in reloaded_model._bqml_model.model_name
     )
     assert reloaded_model.n_clusters == 3
+    assert reloaded_model.init == "RANDOM"
+    assert reloaded_model.distance_type == "EUCLIDEAN"
+    assert reloaded_model.max_iter == 20
+    assert reloaded_model.tol == 0.01
+
+
+def test_cluster_configure_fit_load_params(penguins_df_default_index, dataset_id):
+    model = cluster.KMeans(
+        n_clusters=4,
+        init="random",
+        distance_type="cosine",
+        max_iter=30,
+        tol=0.001,
+    )
+
+    df = penguins_df_default_index.dropna()[
+        [
+            "culmen_length_mm",
+            "culmen_depth_mm",
+            "flipper_length_mm",
+            "sex",
+        ]
+    ]
+
+    # TODO(swast): How should we handle the default index? Currently, we get:
+    # "Column bigframes_index_0_z is not found in the input data to the
+    # EVALUATE function."
+    df = df.reset_index(drop=True)
+
+    model.fit(df)
+
+    # save, load, check n_clusters to ensure configuration was kept
+    reloaded_model = model.to_gbq(
+        f"{dataset_id}.temp_configured_cluster_model", replace=True
+    )
+    assert (
+        f"{dataset_id}.temp_configured_cluster_model"
+        in reloaded_model._bqml_model.model_name
+    )
+    assert reloaded_model.n_clusters == 4
+    assert reloaded_model.init == "RANDOM"
+    assert reloaded_model.distance_type == "COSINE"
+    assert reloaded_model.max_iter == 30
+    assert reloaded_model.tol == 0.001

third_party/bigframes_vendored/sklearn/cluster/_kmeans.py

@@ -31,6 +31,34 @@ class KMeans(_BaseKMeans):
         n_clusters (int, default 8):
             The number of clusters to form as well as the number of centroids to generate.
             Default to 8.
+
+        init ("kmeans++", "random" or "custom", default "kmeans++"):
+            The method of initializing the clusters. Default to "kmeans++"
+
+            kmeas++: Initializes a number of centroids equal to the n_clusters value by using the k-means++ algorithm. Using this approach usually trains a better model than using random cluster initialization.
+            random: Initializes the centroids by randomly selecting a number of data points equal to the n_clusters value from the input data.
+            custom: Initializes the centroids using a provided column of type bool. Uses the rows with a value of True as the initial centroids. You specify the column to use by using the init_col option.
+
+        init_col (str or None, default None):
+            The name of the column to use to initialize the centroids. This column must have a type of bool. If this column contains a value of True for a given row, then uses that row as an initial centroid. The number of True rows in this column must be equal to the value you have specified for the n_clusters option.
+            Only works with init method "custom". Default to None.
+
+        distance_type ("euclidean" or "cosine", default "euclidean"):
+            The type of metric to use to compute the distance between two points.
+            Default to "euclidean".
+
+        max_iter (int, default 20):
+            The maximum number of training iterations, where one iteration represents a single pass of the entire training data. Default to 20.
+
+        tol (float, default 0.01):
+            The minimum relative loss improvement that is necessary to continue training. For example, a value of 0.01 specifies that each iteration must reduce the loss by 1% for training to continue.
+            Default to 0.01.
+
+        warm_start (bool, default False):
+            Determines whether to train a model with new training data, new model options, or both. Unless you explicitly override them, the initial options used to train the model are used for the warm start run.
+            Default to False.
+
+
     """
 
     def fit(