model_garden/notebooks/official/datasets/get_started_bq_datasets.ipynb

In [ ]:

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Get started with BigQuery datasets

Open in Colab

Open in Colab Enterprise

Open in Vertex AI Workbench

View on GitHub

Overview

This tutorial demonstrates how to use Vertex AI in production. This tutorial covers data management: get started with BigQuery datasets.

Learn more about BigQuery datasets and Vertex AI for BigQuery users.

Objective

In this tutorial, you learn how to use BigQuery as a dataset for training with Vertex AI.

This tutorial uses the following Google Cloud ML services:

• Vertex AI datasets
• BigQuery datasets

The steps performed include:

• Create a Vertex AI dataset resource from BigQuery table -- compatible for AutoML training.
• Extract a copy of the dataset from BigQuery to a CSV file in Cloud Storage -- compatible for AutoML or custom training.
• Select rows from a BigQuery dataset into a pandas dataframe -- compatible for custom training.
• Select rows from a BigQuery dataset into a tf.data.Dataset -- compatible for custom training TensorFlow models.
• Select rows from extracted CSV files into a tf.data.Dataset -- compatible for custom training TensorFlow models.
• Create a BigQuery dataset from CSV files.
• Extract data from BigQuery table into a DMatrix -- compatible for custom training XGBoost models.

Recommendations

When doing E2E MLOps on Google Cloud, following are the best practices when dealing with structured (tabular) data in BigQuery:

• For AutoML training:

  • Create a managed dataset with Vertex AI TabularDataset.
  • Use the BigQuery table as the input to the dataset.
  • Specify columns and columns transformations when running the AutoML training pipeline job.

• For custom training:

  • For small datasets:
    • Extract the BigQuery to a pandas dataframe.
    • Preprocess the data in the dataframe.
  • For large datasets:
    • TensorFlow model training:
      • Create a tf.data.Dataset generator from the BigQuery table.
      • Specify the columns for the custrom training.
      • Preprocess the data either:
        • Within the generator (upstream)
        • Within the model (downstream)
    • XGBoost model training:
      • Use BigQuery ML built-in XGBoost training.
      • Alternatively, create a DMatrix generator from CSV files extracted from BigQuery table.
    • PyTorch model training:
      • Extract the BigQuery to a pandas dataframe.
      • Preprocess the data in the dataframe.
      • Create a DataLoader generator from the pandas dataframe.

• Alternatively:

  • Extract the BigQuery table to CSV files.
  • Preprocess the CSV files.
  • Create a tf.data.Dataset generator from the CSV files.

Dataset

The dataset used for this tutorial is the GSOD dataset from BigQuery public datasets. In this version of the dataset you consider the fields year, month and day to predict the value of mean daily temperature (mean_temp).

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage
• BigQuery

Learn about Vertex AI pricing, Cloud Storage pricing and BigQuery pricing and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Get started

Install Vertex AI SDK for Python and other required packages

In [ ]:

! pip3 install --upgrade --quiet google-cloud-aiplatform \
                                 google-cloud-bigquery \
                                 tensorflow \
                                 tensorflow-io \
                                 xgboost \
                                 numpy \
                                 pandas \
                                 pyarrow \
                                 db-dtypes

Restart runtime (Colab only)

To use the newly installed packages, you must restart the runtime on Google Colab.

In [ ]:

import sys

if "google.colab" in sys.modules:

    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)

⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️

Authenticate your notebook environment (Colab only)

Authenticate your environment on Google Colab.

In [ ]:

import sys

if "google.colab" in sys.modules:

    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Vertex AI SDK for Python

To get started using Vertex AI, you must have an existing Google Cloud project and enable the Vertex AI API. Learn more about setting up a project and a development environment.

In [ ]:

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}
LOCATION = "us-central1"  # @param {type:"string"}


from google.cloud import aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION)

Create a Cloud Storage bucket

Create a storage bucket to store intermediate artifacts such as datasets.

In [ ]:

BUCKET_URI = f"gs://your-bucket-name-{PROJECT_ID}-unique"  # @param {type:"string"}

If your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.

In [ ]:

! gsutil mb -l $LOCATION $BUCKET_URI

Import libraries and define constants

In [ ]:

import pandas as pd
import xgboost as xgb
from google.cloud import bigquery

Create BigQuery client

Create the BigQuery client.

In [ ]:

bqclient = bigquery.Client(project=PROJECT_ID)

Location of BigQuery training data.

Now, set the variable IMPORT_FILE to the location of the data table in BigQuery and BQ_TABLE with the table id.

In [ ]:

IMPORT_FILE = "bq://bigquery-public-data.samples.gsod"
BQ_TABLE = "bigquery-public-data.samples.gsod"

Create the dataset

BigQuery input data

Next, create the dataset resource using the create method for the TabularDataset class, which takes the following parameters:

• display_name: The human readable name for the dataset resource.
• bq_source: Import data items from a BigQuery table into the dataset resource.
• labels: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.

Learn more about TabularDataset from BigQuery table.

In [ ]:

dataset = aiplatform.TabularDataset.create(
    display_name="NOAA historical weather data",
    bq_source=[IMPORT_FILE],
    labels={"user_metadata": BUCKET_URI[5:]},
)

label_column = "mean_temp"

print(dataset.resource_name)

Copy the dataset to Cloud Storage

Next, you make a copy of the BigQuery table as a CSV file, to Cloud Storage using the BigQuery extract command.

Learn more about BigQuery command line interface.

In [ ]:

comps = BQ_TABLE.split(".")
BQ_PROJECT_DATASET_TABLE = comps[0] + ":" + comps[1] + "." + comps[2]

! bq --location=us extract --destination_format CSV $BQ_PROJECT_DATASET_TABLE $BUCKET_URI/mydata*.csv

IMPORT_FILES = ! gsutil ls $BUCKET_URI/mydata*.csv

print(IMPORT_FILES)

EXAMPLE_FILE = IMPORT_FILES[0]

! gsutil cat $EXAMPLE_FILE | head

Create the dataset

CSV input data

Next, create the dataset resource using the create method for the TabularDataset class, which takes the following parameters:

• display_name: The human readable name for the dataset resource.
• gcs_source: A list of one or more dataset index files to import the data items into the dataset resource.
• labels: User defined metadata. In this example, you store the location of the Cloud Storage bucket containing the user defined data.

Learn more about TabularDataset from CSV files

In [ ]:

gcs_source = IMPORT_FILES

dataset = aiplatform.TabularDataset.create(
    display_name="NOAA historical weather data",
    gcs_source=gcs_source,
    labels={"user_metadata": BUCKET_URI[5:]},
)


label_column = "mean_temp"

print(dataset.resource_name)

Create a view of the BigQuery dataset

Alternatively, you can create a logical view of a BigQuery dataset that has a subset of the fields.

Learn more about Creating BigQuery views.

In [ ]:

# Set dataset name and view name in BigQuery
BQ_MY_DATASET = "[your-dataset-name]"
BQ_MY_TABLE = "[your-view-name]"

# Otherwise, use the default names
if (
    BQ_MY_DATASET == ""
    or BQ_MY_DATASET is None
    or BQ_MY_DATASET == "[your-dataset-name]"
):
    BQ_MY_DATASET = "mlops_dataset"

if BQ_MY_TABLE == "" or BQ_MY_TABLE is None or BQ_MY_TABLE == "[your-view-name]":
    BQ_MY_TABLE = "mlops_view"

In [ ]:

# Create the resources
! bq --location=US mk -d \
$PROJECT_ID:$BQ_MY_DATASET

sql_script = f'''
CREATE OR REPLACE VIEW `{PROJECT_ID}.{BQ_MY_DATASET}.{BQ_MY_TABLE}`
AS SELECT station_number,year,month,day,mean_temp FROM `{BQ_TABLE}`
'''
print(sql_script)

query = bqclient.query(sql_script)

Read the BigQuery dataset into a pandas dataframe

Next, you read a sample of the dataset into a pandas dataframe using BigQuery list_rows() and to_dataframe() method, as follows:

• list_rows(): Performs a query on the specified table and returns a row iterator to the query results. Optionally specify:

• selected_fields: Subset of fields (columns) to return.

• max_results: The maximum number of rows to return. Same as SQL LIMIT command.

• rows.to_dataframe(): Invokes the row iterator and reads in the data into a pandas dataframe.

Learn more about Loading BigQuery table into a dataframe

In [ ]:

# Download the table.
table = bigquery.TableReference.from_string(BQ_TABLE)

rows = bqclient.list_rows(
    table,
    max_results=500,
    selected_fields=[
        bigquery.SchemaField("station_number", "STRING"),
        bigquery.SchemaField("year", "INTEGER"),
        bigquery.SchemaField("month", "INTEGER"),
        bigquery.SchemaField("day", "INTEGER"),
        bigquery.SchemaField("mean_temp", "FLOAT"),
    ],
)

dataframe = rows.to_dataframe()
print(dataframe.head())

Read the BigQuery dataset into a tf.data.Dataset

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:

• parent: Your project ID.
• project_id: The project ID of the BigQuery table.
• dataset_id: The ID of the BigQuery dataset.
• table_id. The ID of the table within the corresponding BigQuery dataset.
• selected_fields: Subset of fields (columns) to return.
• output_types: The output types of the corresponding fields.
• requested_streams: The number of parallel readers.

Learn more about BigQuery TensorFlow reader.

Learn more about tf.data.Dataset.

In [ ]:

from tensorflow.python.framework import dtypes
from tensorflow_io.bigquery import BigQueryClient

feature_names = "station_number,year,month,day".split(",")

target_name = "mean_temp"


def read_bigquery(project, dataset, table):
    tensorflow_io_bigquery_client = BigQueryClient()
    read_session = tensorflow_io_bigquery_client.read_session(
        parent="projects/" + PROJECT_ID,
        project_id=project,
        dataset_id=dataset,
        table_id=table,
        selected_fields=feature_names + [target_name],
        output_types=[dtypes.string] + [dtypes.int32] * 3 + [dtypes.float32],
        requested_streams=2,
    )

    dataset = read_session.parallel_read_rows()
    return dataset


PROJECT, DATASET, TABLE = IMPORT_FILE.split("/")[-1].split(".")
tf_dataset = read_bigquery(PROJECT, DATASET, TABLE)

print(tf_dataset.take(1))

Read CSV files into a tf.data.Dataset

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:

• filenames: A list of one or more CSV files.
• header: Whether CSV file(s) contain a header.
• select_cols: Subset of fields (columns) to return.
• record_defaults: The output types of the corresponding fields.

Learn more about tf.data CsvDataset

In [ ]:

import tensorflow as tf

feature_names = ["station_number,year,month,day".split(",")]

target_name = "mean_temp"

tf_dataset = tf.data.experimental.CsvDataset(
    filenames=IMPORT_FILES,
    header=True,
    select_cols=feature_names.append(target_name),
    record_defaults=[dtypes.string] + [dtypes.int32] * 3 + [dtypes.float32],
)

print(tf_dataset.take(1))

Create a BigQuery dataset from a pandas dataframe

You can create a BigQuery dataset from a pandas dataframe using the BigQuery create_dataset() and load_table_from_dataframe() methods, as follows:

• create_dataset(): Creates an empty BigQuery dataset, with the following parameters:
• dataset_ref: The DatasetReference created from the dataset_id -- e.g., samples.
• load_table_from_dataframe(): Loads one or more CSV files into a table within the corresponding dataset, with the following parameters:
• dataframe: The dataframe.
• table: The TableReference for the table.
• job_config: Specifications on how to load the dataframe data.

In [ ]:

LOCATION = "us"

SCHEMA = [
    bigquery.SchemaField("station_number", "STRING"),
    bigquery.SchemaField("year", "INTEGER"),
    bigquery.SchemaField("month", "INTEGER"),
    bigquery.SchemaField("day", "INTEGER"),
    bigquery.SchemaField("mean_temp", "FLOAT"),
]


DATASET_ID = "samples"
TABLE_ID = "gsod"


def create_bigquery_dataset(dataset_id):
    dataset = bigquery.Dataset(
        bigquery.dataset.DatasetReference(PROJECT_ID, dataset_id)
    )
    dataset.location = "us"

    try:
        dataset = bqclient.create_dataset(dataset)  # API request
        return True
    except Exception as err:
        print(err)
        if err.code != 409:  # http_client.CONFLICT
            raise
    return False


def load_data_into_bigquery(dataframe, dataset_id, table_id):
    create_bigquery_dataset(dataset_id)
    dataset = bqclient.dataset(dataset_id)
    table = dataset.table(table_id)

    job_config = bigquery.LoadJobConfig(
        # Specify a (partial) schema. All columns are always written to the
        # table. The schema is used to assist in data type definitions.
        schema=[
            bigquery.SchemaField("station_number", "STRING"),
            bigquery.SchemaField("year", "INTEGER"),
            bigquery.SchemaField("month", "INTEGER"),
            bigquery.SchemaField("day", "INTEGER"),
            bigquery.SchemaField("mean_temp", "FLOAT"),
        ],
        # Optionally, set the write disposition. BigQuery appends loaded rows
        # to an existing table by default, but with WRITE_TRUNCATE write
        # disposition it replaces the table with the loaded data.
        write_disposition="WRITE_TRUNCATE",
    )

    NEW_BQ_TABLE = f"{PROJECT_ID}.{dataset_id}.{table_id}"

    job = bqclient.load_table_from_dataframe(
        dataframe, NEW_BQ_TABLE, job_config=job_config
    )  # Make an API request.
    job.result()  # Wait for the job to complete.

    table = bqclient.get_table(NEW_BQ_TABLE)  # Make an API request.
    print(
        "Loaded {} rows and {} columns to {}".format(
            table.num_rows, len(table.schema), NEW_BQ_TABLE
        )
    )


load_data_into_bigquery(dataframe, DATASET_ID, TABLE_ID)

Create a BigQuery dataset from CSV files

You can create a BigQuery dataset from CSV files using the BigQuery create_dataset() and load_table_from_uri() methods, as follows:

• create_dataset(): Creates an empty BigQuery dataset, with the following parameters:
• dataset_ref: The DatasetReference created from the dataset_id -- e.g., samples.
• load_table_from_uri(): Loads one or more CSV files into a table within the corresponding dataset, with the following parameters:
• url: A set of one or more CVS files in Cloud Storage storage.
• table: The TableReference for the table.
• job_config: Specifications on how to load the CSV data.

Learn more about Importing CSV data into BigQuery.

In [ ]:

LOCATION = "us"

CSV_SCHEMA = [
    bigquery.SchemaField("station_number", "STRING"),
    bigquery.SchemaField("wban_number", "STRING"),
    bigquery.SchemaField("year", "INTEGER"),
    bigquery.SchemaField("month", "INTEGER"),
    bigquery.SchemaField("day", "INTEGER"),
    bigquery.SchemaField("mean_temp", "FLOAT"),
    bigquery.SchemaField("num_mean_temp_samples", "INTEGER"),
    bigquery.SchemaField("mean_dew_point", "FLOAT"),
    bigquery.SchemaField("num_mean_dew_point_samples", "INTEGER"),
    bigquery.SchemaField("mean_sealevel_pressure", "FLOAT"),
    bigquery.SchemaField("num_mean_sealevel_pressure_samples", "INTEGER"),
    bigquery.SchemaField("mean_station_pressure", "FLOAT"),
    bigquery.SchemaField("num_mean_station_pressure_samples", "INTEGER"),
    bigquery.SchemaField("mean_visibility", "FLOAT"),
    bigquery.SchemaField("num_mean_visibility_samples", "INTEGER"),
    bigquery.SchemaField("mean_wind_speed", "FLOAT"),
    bigquery.SchemaField("num_mean_wind_speed_samples", "INTEGER"),
    bigquery.SchemaField("max_sustained_wind_speed", "FLOAT"),
    bigquery.SchemaField("max_gust_wind_speed", "FLOAT"),
    bigquery.SchemaField("max_temperature", "FLOAT"),
    bigquery.SchemaField("max_temperature_explicit", "BOOLEAN"),
    bigquery.SchemaField("min_temperature", "FLOAT"),
    bigquery.SchemaField("min_temperature_explicit", "BOOLEAN"),
    bigquery.SchemaField("total_percipitation", "FLOAT"),
    bigquery.SchemaField("snow_depth", "FLOAT"),
    bigquery.SchemaField("fog", "BOOLEAN"),
    bigquery.SchemaField("rain", "BOOLEAN"),
    bigquery.SchemaField("snow", "BOOLEAN"),
    bigquery.SchemaField("hail", "BOOLEAN"),
    bigquery.SchemaField("thunder", "BOOLEAN"),
    bigquery.SchemaField("tornado", "BOOLEAN"),
]


DATASET_ID = "samples"
TABLE_ID = "gsod"


def load_data_into_bigquery(url, dataset_id, table_id):
    create_bigquery_dataset(dataset_id)
    dataset = bqclient.dataset(dataset_id)
    table = dataset.table(table_id)

    job_config = bigquery.LoadJobConfig()
    job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
    job_config.source_format = bigquery.SourceFormat.CSV
    job_config.schema = CSV_SCHEMA
    job_config.skip_leading_rows = 1  # heading

    load_job = bqclient.load_table_from_uri(url, table, job_config=job_config)
    print("Starting job {}".format(load_job.job_id))

    load_job.result()  # Waits for table load to complete.
    print("Job finished.")

    destination_table = bqclient.get_table(table)
    print("Loaded {} rows.".format(destination_table.num_rows))


load_data_into_bigquery(IMPORT_FILES, DATASET_ID, TABLE_ID)

Read BigQuery table into XGboost DMatrix

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.

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.

Learn more about Getting started with built-in XGBoost.

Read pandas table into XGboost DMatrix

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.

In [ ]:

dataframe["station_number"] = pd.to_numeric(dataframe["station_number"])
labels = dataframe["mean_temp"]
data = dataframe.drop(["mean_temp"], axis=1)

dtrain = xgb.DMatrix(data, label=labels)

Read CSV files into XGboost DMatrix

Currently, there is no Cloud Storage support in XGBoost. If you use CSV files for input, you need to download them locally.

In [ ]:

! gsutil cp $EXAMPLE_FILE data.csv

dtrain = xgb.DMatrix("data.csv?format=csv&label_column=4")

Clean up

To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.

Otherwise, you can delete the individual resources you created in this tutorial:

• Vertex AI dataset resource
• Cloud Storage Bucket
• BigQuery dataset

Set delete_storage to True to delete the storage resources used in this notebook.

In [ ]:

import os

# Delete the dataset using the Vertex dataset object
dataset.delete()

# Delete the temporary BigQuery dataset
! bq rm -r -f $PROJECT_ID:$DATASET_ID

delete_storage = False
if delete_storage or os.getenv("IS_TESTING"):
    # Delete the created GCS bucket
    ! gsutil rm -r $BUCKET_URI
    # Delete the created BigQuery datasets
    ! bq rm -r -f $PROJECT_ID:$BQ_MY_DATASET