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 style=\"text-align: center\">\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\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fdatasets%2Fget_started_bq_datasets.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\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\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\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\"><br>\n",
        "      View on GitHub\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": "d1ea81ac77f0"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5d353aa47ac"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9bbcbe73e685"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-bigquery \\\n",
        "                                 tensorflow \\\n",
        "                                 tensorflow-io \\\n",
        "                                 xgboost \\\n",
        "                                 numpy \\\n",
        "                                 pandas \\\n",
        "                                 pyarrow \\\n",
        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "16220914acc5"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "157953ab28f0"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c87a2a5d7e35"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5dccb1c8feb6"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cc7251520a07"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c2fc3d7b6bfa"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "45769dc0c97d"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "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": [
        "**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 $LOCATION $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 pandas as pd\n",
        "import xgboost as xgb\n",
        "from google.cloud import bigquery"
      ]
    },
    {
      "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
}