docs: add a code sample for creating a kmeans model

samples/snippets/conftest.py

@@ -18,6 +18,8 @@
 import pytest
 import test_utils.prefixer
 
+import bigframes.pandas as bpd
+
 prefixer = test_utils.prefixer.Prefixer(
     "python-bigquery-dataframes", "samples/snippets"
 )
@@ -43,6 +45,11 @@ def project_id(bigquery_client: bigquery.Client) -> str:
     return bigquery_client.project
 
 
+@pytest.fixture(autouse=True)
+def reset_session():
+    bpd.reset_session()
+
+
 @pytest.fixture(scope="session")
 def dataset_id(bigquery_client: bigquery.Client, project_id: str) -> Iterator[str]:
     dataset_id = prefixer.create_prefix()

samples/snippets/create_kmeans_model_test.py

@@ -0,0 +1,151 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def test_kmeans_sample(project_id: str):
+    your_gcp_project_id = project_id
+    # [START bigquery_dataframes_bqml_kmeans]
+    import datetime
+
+    import bigframes
+    import bigframes.pandas as bpd
+
+    bigframes.options.bigquery.project = your_gcp_project_id
+    # Compute in the EU multi-region to query the London bicycles dataset.
+    bigframes.options.bigquery.location = "EU"
+
+    # Extract the information you'll need to train the k-means model in this
+    # tutorial. Use the read_gbq function to represent cycle hires
+    # data as a DataFrame.
+    h = bpd.read_gbq(
+        "bigquery-public-data.london_bicycles.cycle_hire",
+        col_order=["start_station_name", "start_station_id", "start_date", "duration"],
+    ).rename(
+        columns={
+            "start_station_name": "station_name",
+            "start_station_id": "station_id",
+        }
+    )
+
+    s = bpd.read_gbq(
+        # Use ST_GEOPOINT and ST_DISTANCE to analyze geographical
+        # data. These functions determine spatial relationships between
+        # geographical features.
+        """
+        SELECT
+        id,
+        ST_DISTANCE(
+            ST_GEOGPOINT(s.longitude, s.latitude),
+            ST_GEOGPOINT(-0.1, 51.5)
+        ) / 1000 AS distance_from_city_center
+        FROM
+        `bigquery-public-data.london_bicycles.cycle_stations` s
+        """
+    )
+
+    # Define Python datetime objects in the UTC timezone for range comparison,
+    # because BigQuery stores timestamp data in the UTC timezone.
+    sample_time = datetime.datetime(2015, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc)
+    sample_time2 = datetime.datetime(2016, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc)
+
+    h = h.loc[(h["start_date"] >= sample_time) & (h["start_date"] <= sample_time2)]
+
+    # Replace each day-of-the-week number with the corresponding "weekday" or
+    # "weekend" label by using the Series.map method.
+    h = h.assign(
+        isweekday=h.start_date.dt.dayofweek.map(
+            {
+                0: "weekday",
+                1: "weekday",
+                2: "weekday",
+                3: "weekday",
+                4: "weekday",
+                5: "weekend",
+                6: "weekend",
+            }
+        )
+    )
+
+    # Supplement each trip in "h" with the station distance information from
+    # "s" by merging the two DataFrames by station ID.
+    merged_df = h.merge(
+        right=s,
+        how="inner",
+        left_on="station_id",
+        right_on="id",
+    )
+
+    # Engineer features to cluster the stations. For each station, find the
+    # average trip duration, number of trips, and distance from city center.
+    stationstats = merged_df.groupby(["station_name", "isweekday"]).agg(
+        {"duration": ["mean", "count"], "distance_from_city_center": "max"}
+    )
+    stationstats.columns = ["duration", "num_trips", "distance_from_city_center"]
+    stationstats = stationstats.sort_values(
+        by="distance_from_city_center", ascending=True
+    ).reset_index()
+
+    # Expected output results: >>> stationstats.head(3)
+    # station_name	isweekday duration  num_trips	distance_from_city_center
+    # Borough Road...	weekday	    1110	    5749	    0.12624
+    # Borough Road...	weekend	    2125	    1774	    0.12624
+    # Webber Street...	weekday	    795	        6517	    0.164021
+    #   3 rows × 5 columns
+
+    # [END bigquery_dataframes_bqml_kmeans]
+
+    # [START bigquery_dataframes_bqml_kmeans_fit]
+
+    from bigframes.ml.cluster import KMeans
+
+    # To determine an optimal number of clusters, construct and fit several
+    # K-Means objects with different values of num_clusters, find the error
+    # measure, and pick the point at which the error measure is at its minimum
+    # value.
+    cluster_model = KMeans(n_clusters=4)
+    cluster_model.fit(stationstats)
+    cluster_model.to_gbq(
+        your_gcp_project_id,  # For example: "bqml_tutorial.sample_model"
+        replace=True,
+    )
+    # [END bigquery_dataframes_bqml_kmeans_fit]
+
+    # [START bigquery_dataframes_bqml_kmeans_predict]
+
+    # Select model you'll use for training. `read_gbq_model` loads model data
+    # from BigQuery, but you could also use the `cluster_model` object from
+    # previous steps.
+    cluster_model = bpd.read_gbq_model(
+        your_gcp_project_id,
+        # For example: "bqml_tutorial.london_station_clusters",
+    )
+
+    # Use 'contains' function to predict which clusters contain the stations
+    # with string "Kennington".
+    stationstats = stationstats.loc[
+        stationstats["station_name"].str.contains("Kennington")
+    ]
+
+    result = cluster_model.predict(stationstats)
+
+    # Expected output results:   >>>results.head(3)
+    # CENTROID...	NEAREST...	station_name  isweekday	 duration num_trips dist...
+    # 	1	[{'CENTROID_ID'...	Borough...	  weekday	  1110	    5749	0.13
+    # 	2	[{'CENTROID_ID'...	Borough...	  weekend	  2125      1774	0.13
+    # 	1	[{'CENTROID_ID'...	Webber...	  weekday	  795	    6517	0.16
+    #   3 rows × 7 columns
+
+    # [END bigquery_dataframes_bqml_kmeans_predict]
+
+    assert result is not None