{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "## Train a linear regression model with BigQuery DataFrames ML\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/googleapis/python-bigquery-dataframes/blob/main/notebooks/ml/bq_dataframes_ml_linear_regression_big.ipynb\">\n",
        "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/colab-logo.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/ml/bq_dataframes_ml_linear_regression_big.ipynb\">\n",
        "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/github-logo.png\" width=\"32\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/ml/bq_dataframes_ml_linear_regression_big.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/ml/bq_dataframes_ml_linear_regression_big.ipynb\">\n",
        "      <img src=\"https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTW1gvOovVlbZAIZylUtf5Iu8-693qS1w5NJw&s\" alt=\"BQ logo\" width=\"35\">\n",
        "      Open in BQ Studio\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24743cf4a1e1"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.11"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates training a linear regression model on Big Data using BigQuery DataFrames ML. BigQuery DataFrames ML provides a provides a scikit-learn-like API for ML powered by the BigQuery engine.\n",
        "\n",
        "Learn more about [BigQuery DataFrames](https://cloud.google.com/python/docs/reference/bigframes/latest)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, we use BigQuery DataFrames to create a linear regression model that predicts the levels of Ozone in the atmosphere.\n",
        "\n",
        "The steps include:\n",
        "\n",
        "- Creating a DataFrame from the BigQuery table.\n",
        "- Cleaning and preparing data using `bigframes.pandas` module.\n",
        "- Creating a linear regression model using `bigframes.ml` module.\n",
        "- Saving the ML model to BigQuery for future use.\n",
        "\n",
        "\n",
        "Let's formally define our problem as: **Train a linear regression model to predict the level of ozone in the atmosphere given the measurements of other constituents and properties of the atmosphere.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08d289fa873f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "In this tutorial we are going to use the [`bigquery-public-data.epa_historical_air_quality`](https://console.cloud.google.com/marketplace/product/epa/historical-air-quality) dataset. To quote the description of the dataset:\n",
        "\n",
        "\"The United States Environmental Protection Agency (EPA) protects both public health and the environment by establishing the standards for national air quality. The EPA provides annual summary data as well as hourly and daily data in the categories of criteria gases, particulates, meteorological, and toxics.\"\n",
        "\n",
        "There are several tables capturing data about the constituents of the atmosphere, see them in the [BigQuery cloud console](https://pantheon.corp.google.com/bigquery?p=bigquery-public-data&d=epa_historical_air_quality&page=dataset). Most tables carry 10's of GBs of data, but that is not an issue with BigQuery DataFrames as the data is efficiently processed at BigQuery without transferring them to the client."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* BigQuery (compute)\n",
        "* BigQuery ML\n",
        "\n",
        "Learn about [BigQuery compute pricing](https://cloud.google.com/bigquery/pricing#analysis_pricing_models)\n",
        "and [BigQuery ML pricing](https://cloud.google.com/bigquery/pricing#bqml),\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "## Installation\n",
        "\n",
        "If you don't have [bigframes](https://pypi.org/project/bigframes/) package already installed, uncomment and execute the following cells to\n",
        "\n",
        "1. Install the package\n",
        "1. Restart the notebook kernel (Jupyter or Colab) to work with the package"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9O0Ka4W2MNF3"
      },
      "outputs": [],
      "source": [
        "# !pip install bigframes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "\n",
        "# import IPython\n",
        "#\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "Complete the tasks in this section to set up your environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oDfTjfACBvJk"
      },
      "source": [
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the BigQuery API](https://console.cloud.google.com/flows/enableapi?apiid=bigquery.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "If you don't know your project ID, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Set the BigQuery location\n",
        "\n",
        "You can also change the `LOCATION` variable used by BigQuery. Learn more about [BigQuery locations](https://cloud.google.com/bigquery/docs/locations#supported_locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eF-Twtc4XGem"
      },
      "outputs": [],
      "source": [
        "LOCATION = \"US\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Set up APIs, IAM permissions and Authentication\n",
        "\n",
        "Follow the instructions at https://cloud.google.com/bigquery/docs/use-bigquery-dataframes#permissions.\n",
        "\n",
        "Depending on your notebook environment, you might have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**Vertex AI Workbench**\n",
        "\n",
        "Do nothing, you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**Local JupyterLab instance**\n",
        "\n",
        "Uncomment and run the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login\n",
        "# ! gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**Colab**\n",
        "\n",
        "Uncomment and run the following cell:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import bigframes.pandas as bpd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Set BigQuery DataFrames options"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NPPMuw2PXGeo"
      },
      "outputs": [],
      "source": [
        "# NOTE: The project option is not required in all environments.\n",
        "# On BigQuery Studio, the project ID is automatically detected.\n",
        "bpd.options.bigquery.project = PROJECT_ID\n",
        "\n",
        "# NOTE: The location option is not required.\n",
        "# It defaults to the location of the first table or query\n",
        "# passed to read_gbq(). For APIs where a location can't be\n",
        "# auto-detected, the location defaults to the \"US\" location.\n",
        "bpd.options.bigquery.location = LOCATION\n",
        "\n",
        "# NOTE: For a machine learning model the order of the data is\n",
        "# not important. So let's relax the ordering_mode to accept\n",
        "# partial ordering. This allows BigQuery DataFrames to run cost\n",
        "# and performance optimized jobs at the BigQuery engine.\n",
        "bpd.options.bigquery.ordering_mode = \"partial\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D21CoOlfFTYI"
      },
      "source": [
        "If you want to reset the location of the created DataFrame or Series objects, reset the session by executing `bpd.close_session()`. After that, you can reuse `bpd.options.bigquery.location` to specify another location."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9EMAqR37AfLS"
      },
      "source": [
        "## Read data in BigQuery tables as DataFrame\n",
        "\n",
        "Let's read the tables in the dataset to construct a BigQuery DataFrames DataFrame. We will combine measurements of various parameters of the atmosphere from multiple tables to represent a consolidated dataframe to use for our model training and prediction. We have daily and hourly versions of the data available, but since we want to create a model that is dynamic so that it can capture the variance throughout the day, we would choose the hourly version.\n",
        "\n",
        "Note that we would use the pandas APIs as we normally would on the BigQuery DataFrames DataFrame, but calculations happen in the BigQuery query engine instead of the local environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "dataset = \"bigquery-public-data.epa_historical_air_quality\"\n",
        "hourly_summary_tables = [\n",
        "    \"co_hourly_summary\",\n",
        "    \"hap_hourly_summary\",\n",
        "    \"no2_hourly_summary\",\n",
        "    \"nonoxnoy_hourly_summary\",\n",
        "    \"o3_hourly_summary\",\n",
        "    \"pm10_hourly_summary\",\n",
        "    \"pm25_frm_hourly_summary\",\n",
        "    \"pm25_nonfrm_hourly_summary\",\n",
        "    \"pm25_speciation_hourly_summary\",\n",
        "    \"pressure_hourly_summary\",\n",
        "    \"rh_and_dp_hourly_summary\",\n",
        "    \"so2_hourly_summary\",\n",
        "    \"temperature_hourly_summary\",\n",
        "    \"voc_hourly_summary\",\n",
        "    \"wind_hourly_summary\",\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's pick index columns - to identify a measurement of the atmospheric parameter, param column - to identify which param the measurement pertains to, and value column - the column containing the measurement itself."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "index_columns = [\"state_name\", \"county_name\", \"site_num\", \"date_local\", \"time_local\"]\n",
        "param_column = \"parameter_name\"\n",
        "value_column = \"sample_measurement\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's observe how much data each table contains:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "for table in hourly_summary_tables:\n",
        "    # get the bigframes global session\n",
        "    bigframes_session = bpd.get_global_session()\n",
        "\n",
        "    # get the bigquery table info\n",
        "    table_info = bigframes_session.bqclient.get_table(f\"{dataset}.{table}\")\n",
        "\n",
        "    # read the table as a dataframe\n",
        "    df = bpd.read_gbq(f\"{dataset}.{table}\")\n",
        "\n",
        "    # print metadata about the table\n",
        "    print(\n",
        "        f\"{table}: \"\n",
        "        f\"{round(table_info.num_bytes/1_000_000_000, 1)} GB, \"\n",
        "        f\"{round(table_info.num_rows/1_000_000, 1)} million rows, \"\n",
        "        f\"{df[param_column].nunique()} params\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's be mindful that the rows in each table may contain duplicates, which may introdude bias in any model trained on the raw data. We will make sure to drop the duplicates when we use the data for model training.\n",
        "\n",
        "Since we want to predict ozone level, we obviously pick the `o3` table. Let's also pick the tables about other gases - `co`, `no2` and `so2`. Let's also pick `pressure` and `temperature` tables as they seem fundamental indicators for the atmosphere. Note that each of these tables capture measurements for a single parameter (i.e. the column `parameter_name` has a single unique value).\n",
        "\n",
        "We are also interested in the nonoxny and wind tables, but they capture multiple parameters (i.e. the column `parameter_name` has a more than one unique values). We will include their measurements in later step, as they require extar processing to separate out the measurements for the individual parameters.\n",
        "\n",
        "We skip the other tables in this exercise for either they have very little or fragmented data or they seem uninteresting for the purpose of predicting ozone levels. You can take this as a separate exercise to train a linear regression model by including those parameters. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's maintain an array of dtaframes, one for each parameter, and eventually combine them into a single dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "params_dfs = []"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's process the tables with single parameter measurements first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EDAaIwHpQCDZ"
      },
      "outputs": [],
      "source": [
        "table_param_dict = {\n",
        "    \"co_hourly_summary\" : \"co\",\n",
        "    \"no2_hourly_summary\" : \"no2\",\n",
        "    \"o3_hourly_summary\" : \"o3\",\n",
        "    \"pressure_hourly_summary\" : \"pressure\",\n",
        "    \"so2_hourly_summary\" : \"so2\",\n",
        "    \"temperature_hourly_summary\" : \"temperature\",\n",
        "}\n",
        "\n",
        "for table, param in table_param_dict.items():\n",
        "    param_df = bpd.read_gbq(\n",
        "        f\"{dataset}.{table}\",\n",
        "        columns=index_columns + [value_column]\n",
        "    )\n",
        "    param_df = param_df\\\n",
        "        .sort_values(index_columns)\\\n",
        "        .drop_duplicates(index_columns)\\\n",
        "        .set_index(index_columns)\\\n",
        "        .rename(columns={value_column : param})\n",
        "    params_dfs.append(param_df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The nonoxnoy table captures measurements for 3 parameters. Let's analyze how many instances of each parameter it contains."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "nonoxnoy_table = f\"{dataset}.nonoxnoy_hourly_summary\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "bpd.read_gbq(nonoxnoy_table, columns=[param_column]).value_counts()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We see that the NOy data is significantly sparse as compared to NO and NOx, so we skip that and include NO and NOx data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "no_df = bpd.read_gbq(\n",
        "    nonoxnoy_table,\n",
        "    columns=index_columns + [value_column],\n",
        "    filters=[(param_column, \"==\", \"Nitric oxide (NO)\")]\n",
        ")\n",
        "no_df = no_df\\\n",
        "    .sort_values(index_columns)\\\n",
        "    .drop_duplicates(index_columns)\\\n",
        "    .set_index(index_columns)\\\n",
        "    .rename(columns={value_column: \"no_\"})\n",
        "params_dfs.append(no_df)\n",
        "\n",
        "nox_df = bpd.read_gbq(\n",
        "    nonoxnoy_table,\n",
        "    columns=index_columns + [value_column],\n",
        "    filters=[(param_column, \"==\", \"Oxides of nitrogen (NOx)\")]\n",
        ")\n",
        "nox_df = nox_df\\\n",
        "    .sort_values(index_columns)\\\n",
        "    .drop_duplicates(index_columns)\\\n",
        "    .set_index(index_columns)\\\n",
        "    .rename(columns={value_column: \"nox\"})\n",
        "params_dfs.append(nox_df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The wind table captures measurements for 2 parameters. Let's analyze how many instances of each parameter it contains."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "wind_table = f\"{dataset}.wind_hourly_summary\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "bpd.read_gbq(wind_table, columns=[param_column]).value_counts()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's include the data for wind speed and wind direction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "wind_speed_df = bpd.read_gbq(\n",
        "    wind_table,\n",
        "    columns=index_columns + [value_column],\n",
        "    filters=[(param_column, \"==\", \"Wind Speed - Resultant\")]\n",
        ")\n",
        "wind_speed_df = wind_speed_df\\\n",
        "    .sort_values(index_columns)\\\n",
        "    .drop_duplicates(index_columns)\\\n",
        "    .set_index(index_columns)\\\n",
        "    .rename(columns={value_column: \"wind_speed\"})\n",
        "params_dfs.append(wind_speed_df)\n",
        "\n",
        "wind_dir_df = bpd.read_gbq(\n",
        "    wind_table,\n",
        "    columns=index_columns + [value_column],\n",
        "    filters=[(param_column, \"==\", \"Wind Direction - Resultant\")]\n",
        ")\n",
        "wind_dir_df = wind_dir_df\\\n",
        "    .sort_values(index_columns)\\\n",
        "    .drop_duplicates(index_columns)\\\n",
        "    .set_index(index_columns)\\\n",
        "    .rename(columns={value_column: \"wind_dir\"})\n",
        "params_dfs.append(wind_dir_df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's observe each individual parameter and number of data points for each parameter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "for param_df in params_dfs:\n",
        "    print(f\"{param_df.columns.values}: {len(param_df)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's combine data from all parameters into a single DataFrame. The measurements for each parameter may not be available for every (state, county, site, date, time) identifier, we will consider only those identifiers for which measurements of all parameters are available. To achieve this we will combine the measurements via \"inner\" join.\n",
        "\n",
        "We will also materialize this combined data via `cache` method for efficient reuse in the subsequent steps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "df = bpd.concat(params_dfs, axis=1, join=\"inner\").cache()\n",
        "df.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rwPLjqW2Ajzh"
      },
      "source": [
        "## Clean and prepare data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's temporarily bring the index columns as dataframe columns for further processing on the index values for the purpose of data preparation.\n",
        "We will reconstruct the index back at the time of the model training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "df = df.reset_index()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Observe the years from which we have consolidated data so far."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "df[\"date_local\"].dt.year.value_counts().sort_index().to_pandas()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this tutorial we would train a model from the past data to predict ozone levels for the future data. Let's define the cut-off year as 2020. We will pretend that the data before 2020 has known ozone levels, and the 2020 onwards the ozone levels are unknown, which we will predict using our model.\n",
        "\n",
        "We should further separate the known data into training and test sets. The model would be trained on the training set and then evaluated on the test set to make sure the model generalizes beyond the training data. We could use [train_test_split](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.model_selection#bigframes_ml_model_selection_train_test_split) method to randomly split the training and test data, but we leave that for you to try out. In this exercise, let's split based on another cutoff year 2017 - the known data before 2017 would be training data and 2017 onwards would be the test data. This way we stay with the idea that the model is trained on past data and then used to predict the future values."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6i6HkFJZa8na"
      },
      "outputs": [],
      "source": [
        "train_data_filter = (df.date_local.dt.year < 2017)\n",
        "test_data_filter = (df.date_local.dt.year >= 2017) & (df.date_local.dt.year < 2020)\n",
        "predict_data_filter = (df.date_local.dt.year >= 2020)\n",
        "\n",
        "df_train = df[train_data_filter].set_index(index_columns)\n",
        "df_test = df[test_data_filter].set_index(index_columns)\n",
        "df_predict = df[predict_data_filter].set_index(index_columns)\n",
        "\n",
        "df_train.shape, df_test.shape, df_predict.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_-0X7NxYK5f"
      },
      "source": [
        "Prepare your feature (or input) columns and the target (or output) column for the purpose of model training and evaluation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YKwCW7Nsavap"
      },
      "outputs": [],
      "source": [
        "X_train = df_train.drop(columns=\"o3\")\n",
        "y_train = df_train[\"o3\"]\n",
        "\n",
        "X_test = df_test.drop(columns=\"o3\")\n",
        "y_test = df_test[\"o3\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prepare the unknown data for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wej78IDUaRW9"
      },
      "outputs": [],
      "source": [
        "X_predict = df_predict.drop(columns=\"o3\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fx4lsNqMorJ-"
      },
      "source": [
        "## Create the linear regression model\n",
        "\n",
        "BigQuery DataFrames ML lets you seamlessly transition from exploring data to creating machine learning models through its scikit-learn-like API, `bigframes.ml`. BigQuery DataFrames ML supports several types of [ML models](https://cloud.google.com/python/docs/reference/bigframes/latest#ml-capabilities).\n",
        "\n",
        "In this notebook, you create a [`LinearRegression`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.linear_model.LinearRegression) model, a type of regression model that generates a continuous value from a linear combination of input features.\n",
        "\n",
        "When you create a model with BigQuery DataFrames ML, it is saved in an internal location and limited to the BigQuery DataFrames session. However, as you'll see in the next section, you can use `to_gbq` to save the model permanently to your BigQuery project."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EloGtMnverFF"
      },
      "source": [
        "### Create the model using `bigframes.ml`\n",
        "\n",
        "Please note that BigQuery DataFrames ML is backed by BigQuery ML, which uses\n",
        "[automatic preprocessing](https://cloud.google.com/bigquery/docs/auto-preprocessing) to encode string values and scale numeric values when you pass the feature columns without transforms.\n",
        "\n",
        "BigQuery ML also [automatically splits the data for training and evaluation](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-glm#data_split_method), although for datasets with less than 500 rows (such as this one), all rows are used for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GskyyUQPowBT"
      },
      "outputs": [],
      "source": [
        "from bigframes.ml.linear_model import LinearRegression\n",
        "\n",
        "model = LinearRegression()\n",
        "\n",
        "model.fit(X_train, y_train)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGjeMPC2caKK"
      },
      "source": [
        "### Score the model\n",
        "\n",
        "Check how the model performs by using the [`score`](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.ml.linear_model.LinearRegression#bigframes_ml_linear_model_LinearRegression_score) method. More information on BigQuery ML model scoring can be found [here](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-evaluate#mlevaluate_output)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kGBJKafpo0dl"
      },
      "outputs": [],
      "source": [
        "# On the training data\n",
        "model.score(X_train, y_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# On the test data\n",
        "model.score(X_test, y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P2lUiZZ_cjri"
      },
      "source": [
        "### Predict using the model\n",
        "\n",
        "Use the model to predict the levels of ozone. The predicted levels are returned in the column `predicted_o3`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bsQ9cmoWo0Ps"
      },
      "outputs": [],
      "source": [
        "df_pred = model.predict(X_predict)\n",
        "df_pred.peek()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GTRdUw-Ro5R1"
      },
      "source": [
        "## Save the model in BigQuery\n",
        "\n",
        "The model is saved locally within this session. You can save the model permanently to BigQuery for use in future sessions, and to make the model sharable with others."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K0mPaoGpcwwy"
      },
      "source": [
        "Create a BigQuery dataset to house the model, adding a name for your dataset as the `DATASET_ID` variable:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZSP7gt13QrQt"
      },
      "outputs": [],
      "source": [
        "DATASET_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "if not DATASET_ID:\n",
        "    raise ValueError(\"Please define the DATASET_ID\")\n",
        "\n",
        "client = bpd.get_global_session().bqclient\n",
        "dataset = client.create_dataset(DATASET_ID, exists_ok=True)\n",
        "print(f\"Dataset {dataset.dataset_id} created.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zqAIWWgJczp-"
      },
      "source": [
        "Save the model using the `to_gbq` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QE_GD4Byo_jb"
      },
      "outputs": [],
      "source": [
        "model.to_gbq(DATASET_ID + \".o3_lr_model\" , replace=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7uHacAy49rT"
      },
      "source": [
        "You can view the saved model in the BigQuery console under the dataset you created in the first step. Run the following cell and follow the link to view your BigQuery console:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qDBoiA_0488Z"
      },
      "outputs": [],
      "source": [
        "print(f'https://console.cloud.google.com/bigquery?ws=!1m5!1m4!5m3!1s{PROJECT_ID}!2s{DATASET_ID}!3so3_lr_model')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G_wjSfXpWTuy"
      },
      "source": [
        "# Summary and next steps\n",
        "\n",
        "You've created a linear regression model using `bigframes.ml`.\n",
        "\n",
        "Learn more about BigQuery DataFrames in the [documentation](https://cloud.google.com/python/docs/reference/bigframes/latest) and find more sample notebooks in the [GitHub repo](https://github.com/googleapis/python-bigquery-dataframes/tree/main/notebooks)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can uncomment the remaining cells and run them to delete the individual resources you created in this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# # Delete the BigQuery dataset and associated ML model\n",
        "# client.delete_dataset(DATASET_ID, delete_contents=True, not_found_ok=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
