notebooks/official/datasets/get_started_bq_datasets.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 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": "title:generic,gcp"
      },
      "source": [
        "# Get started with BigQuery datasets\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/datasets/get_started_bq_datasets.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/datasets/get_started_bq_datasets.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/datasets/get_started_bq_datasets.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",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI in production. This tutorial covers data management: get started with BigQuery datasets.\n",
        "\n",
        "Learn more about [BigQuery Datasets](https://cloud.google.com/bigquery/docs/datasets-intro) and [Vertex AI for BigQuery users](https://cloud.google.com/vertex-ai/docs/beginner/bqml)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage1,get_started_bq"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use `BigQuery` as a dataset for training with `Vertex AI`.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI Datasets`\n",
        "- `BigQuery Datasets`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a Vertex AI `Dataset` resource from `BigQuery` table -- compatible for `AutoML` training.\n",
        "- Extract a copy of the dataset from `BigQuery` to a CSV file in Cloud Storage -- compatible for `AutoML` or custom training.\n",
        "- Select rows from a `BigQuery` dataset into a `pandas` dataframe -- compatible for custom training.\n",
        "- Select rows from a `BigQuery` dataset into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.\n",
        "- Select rows from extracted CSV files into a `tf.data.Dataset` -- compatible for custom training `TensorFlow` models.\n",
        "- Create a `BigQuery` dataset from CSV files.\n",
        "- Extract data from `BigQuery` table into a `DMatrix` -- compatible for custom training `XGBoost` models."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "recommendation:mlops,stage1,tabular,bq"
      },
      "source": [
        "### Recommendations\n",
        "\n",
        "When doing E2E MLOps on Google Cloud, following are the best practices when dealing with structured (tabular) data in BigQuery:\n",
        "\n",
        "- For AutoML training:\n",
        "  - Create a managed dataset with Vertex AI `TabularDataset`.\n",
        "  - Use the BigQuery table as the input to the dataset.\n",
        "  - Specify columns and columns transformations when running the AutoML training pipeline job.\n",
        "\n",
        "\n",
        "- For custom training:\n",
        "  - For small datasets:\n",
        "    - Extract the BigQuery to a pandas dataframe.\n",
        "    - Preprocess the data in the dataframe.\n",
        "  - For large datasets:\n",
        "    - TensorFlow model training:\n",
        "      - Create a tf.data.Dataset generator from the BigQuery table.\n",
        "      - Specify the columns for the custrom training.\n",
        "      - Preprocess the data either:\n",
        "        - Within the generator (upstream)\n",
        "        - Within the model (downstream)\n",
        "    - XGBoost model training:\n",
        "      - Use BigQuery ML built-in XGBoost training.\n",
        "      - Alternatively, create a DMatrix generator from CSV files extracted from BigQuery table.\n",
        "    - PyTorch model training:\n",
        "        - Extract the BigQuery to a pandas dataframe.\n",
        "        - Preprocess the data in the dataframe.\n",
        "        - Create a DataLoader generator from the pandas dataframe.\n",
        "\n",
        "\n",
        "- Alternatively:\n",
        "    - Extract the BigQuery table to CSV files.\n",
        "    - Preprocess the CSV files.\n",
        "    - Create a tf.data.Dataset generator from the CSV files."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:gsod,lrg"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the GSOD dataset from [BigQuery public datasets](https://cloud.google.com/bigquery/public-data). In this version of the dataset you consider the fields year, month and day to predict the value of mean daily temperature (mean_temp)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9e483012a752"
      },
      "source": [
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "- BigQuery\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and [BigQuery pricing](https://cloud.google.com/bigquery/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_mlops"
      },
      "source": [
        "## Installations\n",
        "\n",
        "Install the following packages to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_mlops"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-bigquery \\\n",
        "                                 tensorflow \\\n",
        "                                 tensorflow-io==0.18 \\\n",
        "                                 xgboost \\\n",
        "                                 numpy \\\n",
        "                                 pandas \\\n",
        "                                 pyarrow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yfEglUHQk9S3"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### 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": "set_project_id"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcp_authenticate"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce6043da7b33"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0367eac06a10"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ad4dbb4a61"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:mbsdk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import google.cloud.aiplatform as aiplatform\n",
        "import pandas as pd\n",
        "import xgboost as xgb\n",
        "from google.cloud import bigquery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,region"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk,region"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_bq"
      },
      "source": [
        "### Create BigQuery client\n",
        "\n",
        "Create the BigQuery client."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_bq"
      },
      "outputs": [],
      "source": [
        "bqclient = bigquery.Client(project=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "import_file:u_dataset,bq"
      },
      "source": [
        "#### Location of BigQuery training data.\n",
        "\n",
        "Now, set the variable `IMPORT_FILE` to the location of the data table in BigQuery and `BQ_TABLE` with the table id."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_file:gsod,bq,lrg"
      },
      "outputs": [],
      "source": [
        "IMPORT_FILE = \"bq://bigquery-public-data.samples.gsod\"\n",
        "BQ_TABLE = \"bigquery-public-data.samples.gsod\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_dataset:tabular,bq,lrg"
      },
      "source": [
        "### Create the Dataset\n",
        "\n",
        "#### BigQuery input data\n",
        "\n",
        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the `Dataset` resource.\n",
        "- `bq_source`: Import data items from a BigQuery table into the `Dataset` resource.\n",
        "- `labels`: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.\n",
        "\n",
        "Learn more about [TabularDataset from BigQuery table](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_bigquery_sample-python)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_dataset:tabular,bq,lrg"
      },
      "outputs": [],
      "source": [
        "dataset = aiplatform.TabularDataset.create(\n",
        "    display_name=\"NOAA historical weather data\",\n",
        "    bq_source=[IMPORT_FILE],\n",
        "    labels={\"user_metadata\": BUCKET_URI[5:]},\n",
        ")\n",
        "\n",
        "label_column = \"mean_temp\"\n",
        "\n",
        "print(dataset.resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bq_extract"
      },
      "source": [
        "### Copy the dataset to Cloud Storage\n",
        "\n",
        "Next, you make a copy of the BigQuery table as a CSV file, to Cloud Storage using the BigQuery extract command.\n",
        "\n",
        "Learn more about [BigQuery command line interface](https://cloud.google.com/bigquery/docs/reference/bq-cli-reference)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bq_extract"
      },
      "outputs": [],
      "source": [
        "comps = BQ_TABLE.split(\".\")\n",
        "BQ_PROJECT_DATASET_TABLE = comps[0] + \":\" + comps[1] + \".\" + comps[2]\n",
        "\n",
        "! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_URI/mydata*.csv\n",
        "\n",
        "IMPORT_FILES = ! gsutil ls $BUCKET_URI/mydata*.csv\n",
        "\n",
        "print(IMPORT_FILES)\n",
        "\n",
        "EXAMPLE_FILE = IMPORT_FILES[0]\n",
        "\n",
        "! gsutil cat $EXAMPLE_FILE | head"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_dataset:tabular,lrg"
      },
      "source": [
        "### Create the Dataset\n",
        "\n",
        "#### CSV input data\n",
        "\n",
        "Next, create the `Dataset` resource using the `create` method for the `TabularDataset` class, which takes the following parameters:\n",
        "\n",
        "- `display_name`: The human readable name for the `Dataset` resource.\n",
        "- `gcs_source`: A list of one or more dataset index files to import the data items into the `Dataset` resource.\n",
        "- `labels`: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.\n",
        "\n",
        "Learn more about [TabularDataset from CSV files](https://cloud.google.com/vertex-ai/docs/datasets/create-dataset-api#aiplatform_create_dataset_tabular_gcs_sample-python)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_dataset:tabular,lrg"
      },
      "outputs": [],
      "source": [
        "gcs_source = IMPORT_FILES\n",
        "\n",
        "dataset = aiplatform.TabularDataset.create(\n",
        "    display_name=\"NOAA historical weather data\",\n",
        "    gcs_source=gcs_source,\n",
        "    labels={\"user_metadata\": BUCKET_URI[5:]},\n",
        ")\n",
        "\n",
        "\n",
        "label_column = \"mean_temp\"\n",
        "\n",
        "print(dataset.resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bq_view"
      },
      "source": [
        "### Create a view of the BigQuery dataset\n",
        "\n",
        "Alternatively, you can create a logical view of a BigQuery dataset that has a subset of the fields.\n",
        "\n",
        "Learn more about [Creating BigQuery views](https://cloud.google.com/bigquery/docs/views)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7dc142433e50"
      },
      "outputs": [],
      "source": [
        "# Set dataset name and view name in BigQuery\n",
        "BQ_MY_DATASET = \"[your-dataset-name]\"\n",
        "BQ_MY_TABLE = \"[your-view-name]\"\n",
        "\n",
        "# Otherwise, use the default names\n",
        "if (\n",
        "    BQ_MY_DATASET == \"\"\n",
        "    or BQ_MY_DATASET is None\n",
        "    or BQ_MY_DATASET == \"[your-dataset-name]\"\n",
        "):\n",
        "    BQ_MY_DATASET = \"mlops_dataset\"\n",
        "\n",
        "if BQ_MY_TABLE == \"\" or BQ_MY_TABLE is None or BQ_MY_TABLE == \"[your-view-name]\":\n",
        "    BQ_MY_TABLE = \"mlops_view\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bq_view"
      },
      "outputs": [],
      "source": [
        "# Create the resources\n",
        "! bq --location=US mk -d \\\n",
        "$PROJECT_ID:$BQ_MY_DATASET\n",
        "\n",
        "sql_script = f'''\n",
        "CREATE OR REPLACE VIEW `{PROJECT_ID}.{BQ_MY_DATASET}.{BQ_MY_TABLE}`\n",
        "AS SELECT station_number,year,month,day,mean_temp FROM `{BQ_TABLE}`\n",
        "'''\n",
        "print(sql_script)\n",
        "\n",
        "query = bqclient.query(sql_script)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bq_to_dataframe:gsod"
      },
      "source": [
        "### Read the BigQuery dataset into a pandas dataframe\n",
        "\n",
        "Next, you read a sample of the dataset into a pandas dataframe using BigQuery `list_rows()` and `to_dataframe()` method, as follows:\n",
        "\n",
        "- `list_rows()`: Performs a query on the specified table and returns a row iterator to the query results. Optionally specify:\n",
        " - `selected_fields`: Subset of fields (columns) to return.\n",
        " - `max_results`: The maximum number of rows to return. Same as SQL LIMIT command.\n",
        "\n",
        "\n",
        "- `rows.to_dataframe()`: Invokes the row iterator and reads in the data into a pandas dataframe.\n",
        "\n",
        "Learn more about [Loading BigQuery table into a dataframe](https://cloud.google.com/bigquery/docs/bigquery-storage-python-pandas)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bq_to_dataframe:gsod"
      },
      "outputs": [],
      "source": [
        "# Download the table.\n",
        "table = bigquery.TableReference.from_string(BQ_TABLE)\n",
        "\n",
        "rows = bqclient.list_rows(\n",
        "    table,\n",
        "    max_results=500,\n",
        "    selected_fields=[\n",
        "        bigquery.SchemaField(\"station_number\", \"STRING\"),\n",
        "        bigquery.SchemaField(\"year\", \"INTEGER\"),\n",
        "        bigquery.SchemaField(\"month\", \"INTEGER\"),\n",
        "        bigquery.SchemaField(\"day\", \"INTEGER\"),\n",
        "        bigquery.SchemaField(\"mean_temp\", \"FLOAT\"),\n",
        "    ],\n",
        ")\n",
        "\n",
        "dataframe = rows.to_dataframe()\n",
        "print(dataframe.head())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bq_to_dataset:gsod"
      },
      "source": [
        "### Read the BigQuery dataset into a tf.data.Dataset\n",
        "\n",
        "Next, you read a sample of the dataset into a tf.data.Dataset using TensorFlow IO `BigQueryClient()` and `read_session()` method, with the following parameters:\n",
        "\n",
        "- `parent`: Your project ID.\n",
        "- `project_id`: The project ID of the BigQuery table.\n",
        "- `dataset_id`: The ID of the BigQuery dataset.\n",
        "- `table_id`. The ID of the table within the corresponding BigQuery dataset.\n",
        "- `selected_fields`: Subset of fields (columns) to return.\n",
        "- `output_types`: The output types of the corresponding fields.\n",
        "- `requested_streams`: The number of parallel readers.\n",
        "\n",
        "Learn more about [BigQuery TensorFlow reader](https://www.tensorflow.org/io/tutorials/bigquery).\n",
        "\n",
        "Learn more about [tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bq_to_dataset:gsod"
      },
      "outputs": [],
      "source": [
        "from tensorflow.python.framework import dtypes\n",
        "from tensorflow_io.bigquery import BigQueryClient\n",
        "\n",
        "feature_names = \"station_number,year,month,day\".split(\",\")\n",
        "\n",
        "target_name = \"mean_temp\"\n",
        "\n",
        "\n",
        "def read_bigquery(project, dataset, table):\n",
        "    tensorflow_io_bigquery_client = BigQueryClient()\n",
        "    read_session = tensorflow_io_bigquery_client.read_session(\n",
        "        parent=\"projects/\" + PROJECT_ID,\n",
        "        project_id=project,\n",
        "        dataset_id=dataset,\n",
        "        table_id=table,\n",
        "        selected_fields=feature_names + [target_name],\n",
        "        output_types=[dtypes.string] + [dtypes.int32] * 3 + [dtypes.float32],\n",
        "        requested_streams=2,\n",
        "    )\n",
        "\n",
        "    dataset = read_session.parallel_read_rows()\n",
        "    return dataset\n",
        "\n",
        "\n",
        "PROJECT, DATASET, TABLE = IMPORT_FILE.split(\"/\")[-1].split(\".\")\n",
        "tf_dataset = read_bigquery(PROJECT, DATASET, TABLE)\n",
        "\n",
        "print(tf_dataset.take(1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "csv_to_dataset:gsod"
      },
      "source": [
        "### Read CSV files into a tf.data.Dataset\n",
        "\n",
        "Alternatively, when your data is in CSV files, you can load the dataset into a tf.data.Dataset using `tf.data.experimental.CsvDataset`, with the following parameters:\n",
        "\n",
        "- `filenames`: A list of one or more CSV files.\n",
        "- `header`: Whether CSV file(s) contain a header.\n",
        "- `select_cols`: Subset of fields (columns) to return.\n",
        "- `record_defaults`: The output types of the corresponding fields.\n",
        "\n",
        "Learn more about [tf.data CsvDataset](https://www.tensorflow.org/api_docs/python/tf/data/experimental/CsvDataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csv_to_dataset:gsod"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "feature_names = [\"station_number,year,month,day\".split(\",\")]\n",
        "\n",
        "target_name = \"mean_temp\"\n",
        "\n",
        "tf_dataset = tf.data.experimental.CsvDataset(\n",
        "    filenames=IMPORT_FILES,\n",
        "    header=True,\n",
        "    select_cols=feature_names.append(target_name),\n",
        "    record_defaults=[dtypes.string] + [dtypes.int32] * 3 + [dtypes.float32],\n",
        ")\n",
        "\n",
        "print(tf_dataset.take(1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataframe_to_bq"
      },
      "source": [
        "### Create a BigQuery dataset from a pandas dataframe\n",
        "\n",
        "You can create a BigQuery dataset from a pandas dataframe using the BigQuery `create_dataset()` and `load_table_from_dataframe()` methods, as follows:\n",
        "\n",
        "- `create_dataset()`: Creates an empty BigQuery dataset, with the following parameters:\n",
        " - `dataset_ref`: The `DatasetReference` created from the dataset_id -- e.g., samples.\n",
        "- `load_table_from_dataframe()`: Loads one or more CSV files into a table within the corresponding dataset, with the following parameters:\n",
        " - `dataframe`: The dataframe.\n",
        " - `table`: The `TableReference` for the table.\n",
        " - `job_config`: Specifications on how to load the dataframe data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dataframe_to_bq"
      },
      "outputs": [],
      "source": [
        "LOCATION = \"us\"\n",
        "\n",
        "SCHEMA = [\n",
        "    bigquery.SchemaField(\"station_number\", \"STRING\"),\n",
        "    bigquery.SchemaField(\"year\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"month\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"day\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_temp\", \"FLOAT\"),\n",
        "]\n",
        "\n",
        "\n",
        "DATASET_ID = \"samples\"\n",
        "TABLE_ID = \"gsod\"\n",
        "\n",
        "\n",
        "def create_bigquery_dataset(dataset_id):\n",
        "    dataset = bigquery.Dataset(\n",
        "        bigquery.dataset.DatasetReference(PROJECT_ID, dataset_id)\n",
        "    )\n",
        "    dataset.location = \"us\"\n",
        "\n",
        "    try:\n",
        "        dataset = bqclient.create_dataset(dataset)  # API request\n",
        "        return True\n",
        "    except Exception as err:\n",
        "        print(err)\n",
        "        if err.code != 409:  # http_client.CONFLICT\n",
        "            raise\n",
        "    return False\n",
        "\n",
        "\n",
        "def load_data_into_bigquery(dataframe, dataset_id, table_id):\n",
        "    create_bigquery_dataset(dataset_id)\n",
        "    dataset = bqclient.dataset(dataset_id)\n",
        "    table = dataset.table(table_id)\n",
        "\n",
        "    job_config = bigquery.LoadJobConfig(\n",
        "        # Specify a (partial) schema. All columns are always written to the\n",
        "        # table. The schema is used to assist in data type definitions.\n",
        "        schema=[\n",
        "            bigquery.SchemaField(\"station_number\", \"STRING\"),\n",
        "            bigquery.SchemaField(\"year\", \"INTEGER\"),\n",
        "            bigquery.SchemaField(\"month\", \"INTEGER\"),\n",
        "            bigquery.SchemaField(\"day\", \"INTEGER\"),\n",
        "            bigquery.SchemaField(\"mean_temp\", \"FLOAT\"),\n",
        "        ],\n",
        "        # Optionally, set the write disposition. BigQuery appends loaded rows\n",
        "        # to an existing table by default, but with WRITE_TRUNCATE write\n",
        "        # disposition it replaces the table with the loaded data.\n",
        "        write_disposition=\"WRITE_TRUNCATE\",\n",
        "    )\n",
        "\n",
        "    NEW_BQ_TABLE = f\"{PROJECT_ID}.{dataset_id}.{table_id}\"\n",
        "\n",
        "    job = bqclient.load_table_from_dataframe(\n",
        "        dataframe, NEW_BQ_TABLE, job_config=job_config\n",
        "    )  # Make an API request.\n",
        "    job.result()  # Wait for the job to complete.\n",
        "\n",
        "    table = bqclient.get_table(NEW_BQ_TABLE)  # Make an API request.\n",
        "    print(\n",
        "        \"Loaded {} rows and {} columns to {}\".format(\n",
        "            table.num_rows, len(table.schema), NEW_BQ_TABLE\n",
        "        )\n",
        "    )\n",
        "\n",
        "\n",
        "load_data_into_bigquery(dataframe, DATASET_ID, TABLE_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "csv_to_bq"
      },
      "source": [
        "### Create a BigQuery dataset from CSV files\n",
        "\n",
        "You can create a BigQuery dataset from CSV files using the BigQuery `create_dataset()` and `load_table_from_uri()` methods, as follows:\n",
        "\n",
        "- `create_dataset()`: Creates an empty BigQuery dataset, with the following parameters:\n",
        " - `dataset_ref`: The `DatasetReference` created from the dataset_id -- e.g., samples.\n",
        "- `load_table_from_uri()`: Loads one or more CSV files into a table within the corresponding dataset, with the following parameters:\n",
        " - `url`: A set of one or more CVS files in Cloud Storage storage.\n",
        " - `table`: The `TableReference` for the table.\n",
        " - `job_config`: Specifications on how to load the CSV data.\n",
        "\n",
        "Learn more about [Importing CSV data into BigQuery](https://www.tensorflow.org/io/tutorials/bigquery#import_census_data_into_bigquery)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csv_to_bq"
      },
      "outputs": [],
      "source": [
        "LOCATION = \"us\"\n",
        "\n",
        "CSV_SCHEMA = [\n",
        "    bigquery.SchemaField(\"station_number\", \"STRING\"),\n",
        "    bigquery.SchemaField(\"wban_number\", \"STRING\"),\n",
        "    bigquery.SchemaField(\"year\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"month\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"day\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_temp\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_temp_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_dew_point\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_dew_point_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_sealevel_pressure\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_sealevel_pressure_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_station_pressure\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_station_pressure_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_visibility\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_visibility_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"mean_wind_speed\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"num_mean_wind_speed_samples\", \"INTEGER\"),\n",
        "    bigquery.SchemaField(\"max_sustained_wind_speed\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"max_gust_wind_speed\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"max_temperature\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"max_temperature_explicit\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"min_temperature\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"min_temperature_explicit\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"total_percipitation\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"snow_depth\", \"FLOAT\"),\n",
        "    bigquery.SchemaField(\"fog\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"rain\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"snow\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"hail\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"thunder\", \"BOOLEAN\"),\n",
        "    bigquery.SchemaField(\"tornado\", \"BOOLEAN\"),\n",
        "]\n",
        "\n",
        "\n",
        "DATASET_ID = \"samples\"\n",
        "TABLE_ID = \"gsod\"\n",
        "\n",
        "\n",
        "def load_data_into_bigquery(url, dataset_id, table_id):\n",
        "    create_bigquery_dataset(dataset_id)\n",
        "    dataset = bqclient.dataset(dataset_id)\n",
        "    table = dataset.table(table_id)\n",
        "\n",
        "    job_config = bigquery.LoadJobConfig()\n",
        "    job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE\n",
        "    job_config.source_format = bigquery.SourceFormat.CSV\n",
        "    job_config.schema = CSV_SCHEMA\n",
        "    job_config.skip_leading_rows = 1  # heading\n",
        "\n",
        "    load_job = bqclient.load_table_from_uri(url, table, job_config=job_config)\n",
        "    print(\"Starting job {}\".format(load_job.job_id))\n",
        "\n",
        "    load_job.result()  # Waits for table load to complete.\n",
        "    print(\"Job finished.\")\n",
        "\n",
        "    destination_table = bqclient.get_table(table)\n",
        "    print(\"Loaded {} rows.\".format(destination_table.num_rows))\n",
        "\n",
        "\n",
        "load_data_into_bigquery(IMPORT_FILES, DATASET_ID, TABLE_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bq_to_xgboost"
      },
      "source": [
        "### Read BigQuery table into XGboost DMatrix\n",
        "\n",
        "Currently, there is no direct data feeding connector between BigQuery and the open source XGBoost. The BigQuery ML service has a built-in XGBoost training module.\n",
        "\n",
        "Alernatively, you extract the data either as a pandas dataframe or as CSV files. The extracted data is then given as an input to a `DMatrix` object when training the model.\n",
        "\n",
        "Learn more about [Getting started with built-in XGBoost](https://cloud.google.com/ai-platform/training/docs/algorithms/xgboost-start)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pandas_to_xgboost:gsod"
      },
      "source": [
        "### Read pandas table into XGboost DMatrix\n",
        "\n",
        "Next, you load the pandas dataframe into a `DMatrix` object. XGBoost does not support non-numeric inputs. Any column that is categorical need to be one-hot encoded prior to loading the dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pandas_to_xgboost:gsod"
      },
      "outputs": [],
      "source": [
        "dataframe[\"station_number\"] = pd.to_numeric(dataframe[\"station_number\"])\n",
        "labels = dataframe[\"mean_temp\"]\n",
        "data = dataframe.drop([\"mean_temp\"], axis=1)\n",
        "\n",
        "dtrain = xgb.DMatrix(data, label=labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "csv_to_xgboost:gsod"
      },
      "source": [
        "### Read CSV files into XGboost DMatrix\n",
        "\n",
        "Currently, there is no Cloud Storage support in XGBoost. If you use CSV files for input, you need to download them locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csv_to_xgboost:gsod"
      },
      "outputs": [],
      "source": [
        "! gsutil cp $EXAMPLE_FILE data.csv\n",
        "\n",
        "dtrain = xgb.DMatrix(\"data.csv?format=csv&label_column=4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:mbsdk"
      },
      "source": [
        "# Clean 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 delete the individual resources you created in this tutorial:\n",
        "\n",
        "- Vertex AI Dataset resource\n",
        "- Cloud Storage Bucket\n",
        "- BigQuery Dataset\n",
        "\n",
        "Set `delete_storage` to _True_ to delete the storage resources used in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "47ad926d84e8"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Delete the dataset using the Vertex dataset object\n",
        "dataset.delete()\n",
        "\n",
        "# Delete the temporary BigQuery dataset\n",
        "! bq rm -r -f $PROJECT_ID:$DATASET_ID\n",
        "\n",
        "delete_storage = False\n",
        "if delete_storage or os.getenv(\"IS_TESTING\"):\n",
        "    # Delete the created GCS bucket\n",
        "    ! gsutil rm -r $BUCKET_URI\n",
        "    # Delete the created BigQuery datasets\n",
        "    ! bq rm -r -f $PROJECT_ID:$BQ_MY_DATASET"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_bq_datasets.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}