01-dataset-management.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e8ba7f4f",
   "metadata": {},
   "source": [
    "# 01 - Data Analysis and Preparation\n",
    "\n",
    "This notebook covers the following tasks:\n",
    "\n",
    "1. Perform exploratory data analysis and visualization.\n",
    "2. Prepare the data for the ML task in BigQuery.\n",
    "3. Generate and fix a ` TFDV schema` for the source data.\n",
    "4. Create a `Vertex Dataset resource` dataset.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b481a247",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "The [Chicago Taxi Trips](https://pantheon.corp.google.com/marketplace/details/city-of-chicago-public-data/chicago-taxi-trips) dataset is one of [public datasets hosted with BigQuery](https://cloud.google.com/bigquery/public-data/), which includes taxi trips from 2013 to the present, reported to the City of Chicago in its role as a regulatory agency. The `taxi_trips` table size is 70.72 GB and includes more than 195 million records. The dataset includes information about the trips, like pickup and dropoff datetime and location, passengers count, miles travelled, and trip toll. \n",
    "\n",
    "The ML task is to predict whether a given trip will result in a tip > 20%."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fedd0ac",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b25967c9",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2cf3f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_data_validation as tfdv\n",
    "from google.cloud import bigquery\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from google.cloud import aiplatform as vertex_ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd0ee37",
   "metadata": {},
   "source": [
    "### Setup Google Cloud project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7e4712e",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = '[your-project-id]' # Change to your project id.\n",
    "REGION = 'us-central1' # Change to your region.\n",
    "\n",
    "if PROJECT == \"\" or PROJECT is None or PROJECT == \"[your-project-id]\":\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT = shell_output[0]\n",
    "    \n",
    "print(\"Project ID:\", PROJECT)\n",
    "print(\"Region:\", REGION)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "384a817b",
   "metadata": {},
   "source": [
    "### Set configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71d0738",
   "metadata": {},
   "outputs": [],
   "source": [
    "BQ_DATASET_NAME = 'playground_us' # Change to your BQ dataset name.\n",
    "BQ_TABLE_NAME = 'chicago_taxitrips_prep'\n",
    "BQ_LOCATION = 'US'\n",
    "\n",
    "DATASET_DISPLAY_NAME = 'chicago-taxi-tips'\n",
    "\n",
    "RAW_SCHEMA_DIR = 'src/raw_schema'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39395512",
   "metadata": {},
   "source": [
    "## 1. Explore the data in BigQuery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e4300d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bigquery data\n",
    "\n",
    "SELECT \n",
    "    CAST(EXTRACT(DAYOFWEEK FROM trip_start_timestamp) AS string) AS trip_dayofweek, \n",
    "    FORMAT_DATE('%A',cast(trip_start_timestamp as date)) AS trip_dayname,\n",
    "    COUNT(*) as trip_count,\n",
    "FROM `bigquery-public-data.chicago_taxi_trips.taxi_trips`\n",
    "WHERE\n",
    "    EXTRACT(YEAR FROM trip_start_timestamp) = 2015 \n",
    "GROUP BY\n",
    "    trip_dayofweek,\n",
    "    trip_dayname\n",
    "ORDER BY\n",
    "    trip_dayofweek\n",
    ";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699804c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f2447e",
   "metadata": {},
   "outputs": [],
   "source": [
    "data.plot(kind='bar', x='trip_dayname', y='trip_count')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7782c69",
   "metadata": {},
   "source": [
    "## 2. Create data for the ML task\n",
    "\n",
    "We add a `ML_use` column for pre-splitting the data, where 80% of the datsa items are set to `UNASSIGNED` while the other 20% is set to `TEST`.\n",
    "\n",
    "This column is used during training to split the dataset for training and test.\n",
    "\n",
    "In the training phase, the `UNASSIGNED` are split into `train` and `eval`. The `TEST` split is will be used for the final model validation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7987d132",
   "metadata": {},
   "source": [
    "### Create destination BigQuery dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a7f6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!bq --location=$BQ_LOCATION mk -d \\\n",
    "$PROJECT:$BQ_DATASET_NAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "832f99ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 1000000\n",
    "year = 2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b19789f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_script = '''\n",
    "CREATE OR REPLACE TABLE `@PROJECT.@DATASET.@TABLE` \n",
    "AS (\n",
    "    WITH\n",
    "      taxitrips AS (\n",
    "      SELECT\n",
    "        trip_start_timestamp,\n",
    "        trip_seconds,\n",
    "        trip_miles,\n",
    "        payment_type,\n",
    "        pickup_longitude,\n",
    "        pickup_latitude,\n",
    "        dropoff_longitude,\n",
    "        dropoff_latitude,\n",
    "        tips,\n",
    "        fare\n",
    "      FROM\n",
    "        `bigquery-public-data.chicago_taxi_trips.taxi_trips`\n",
    "      WHERE 1=1 \n",
    "      AND pickup_longitude IS NOT NULL\n",
    "      AND pickup_latitude IS NOT NULL\n",
    "      AND dropoff_longitude IS NOT NULL\n",
    "      AND dropoff_latitude IS NOT NULL\n",
    "      AND trip_miles > 0\n",
    "      AND trip_seconds > 0\n",
    "      AND fare > 0\n",
    "      AND EXTRACT(YEAR FROM trip_start_timestamp) = @YEAR\n",
    "    )\n",
    "\n",
    "    SELECT\n",
    "      trip_start_timestamp,\n",
    "      EXTRACT(MONTH from trip_start_timestamp) as trip_month,\n",
    "      EXTRACT(DAY from trip_start_timestamp) as trip_day,\n",
    "      EXTRACT(DAYOFWEEK from trip_start_timestamp) as trip_day_of_week,\n",
    "      EXTRACT(HOUR from trip_start_timestamp) as trip_hour,\n",
    "      trip_seconds,\n",
    "      trip_miles,\n",
    "      payment_type,\n",
    "      ST_AsText(\n",
    "          ST_SnapToGrid(ST_GeogPoint(pickup_longitude, pickup_latitude), 0.1)\n",
    "      ) AS pickup_grid,\n",
    "      ST_AsText(\n",
    "          ST_SnapToGrid(ST_GeogPoint(dropoff_longitude, dropoff_latitude), 0.1)\n",
    "      ) AS dropoff_grid,\n",
    "      ST_Distance(\n",
    "          ST_GeogPoint(pickup_longitude, pickup_latitude), \n",
    "          ST_GeogPoint(dropoff_longitude, dropoff_latitude)\n",
    "      ) AS euclidean,\n",
    "      CONCAT(\n",
    "          ST_AsText(ST_SnapToGrid(ST_GeogPoint(pickup_longitude,\n",
    "              pickup_latitude), 0.1)), \n",
    "          ST_AsText(ST_SnapToGrid(ST_GeogPoint(dropoff_longitude,\n",
    "              dropoff_latitude), 0.1))\n",
    "      ) AS loc_cross,\n",
    "      IF((tips/fare >= 0.2), 1, 0) AS tip_bin,\n",
    "      IF(RAND() <= 0.8, 'UNASSIGNED', 'TEST') AS ML_use\n",
    "    FROM\n",
    "      taxitrips\n",
    "    LIMIT @LIMIT\n",
    ")\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f1d2837",
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_script = sql_script.replace(\n",
    "    '@PROJECT', PROJECT).replace(\n",
    "    '@DATASET', BQ_DATASET_NAME).replace(\n",
    "    '@TABLE', BQ_TABLE_NAME).replace(\n",
    "    '@YEAR', str(year)).replace(\n",
    "    '@LIMIT', str(sample_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27acb6b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sql_script)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f31e8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "bq_client = bigquery.Client(project=PROJECT, location=BQ_LOCATION)\n",
    "job = bq_client.query(sql_script)\n",
    "_ = job.result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b32eea6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bigquery --project {PROJECT}\n",
    "\n",
    "SELECT ML_use, COUNT(*)\n",
    "FROM playground_us.chicago_taxitrips_prep # Change to your BQ dataset and table names.\n",
    "GROUP BY ML_use"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fcc75fc",
   "metadata": {},
   "source": [
    "### Load a sample data to a Pandas DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f252a846",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bigquery sample_data --project {PROJECT}\n",
    "\n",
    "SELECT * EXCEPT (trip_start_timestamp, ML_use)\n",
    "FROM playground_us.chicago_taxitrips_prep # Change to your BQ dataset and table names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798809d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_data.head().T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23595838",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_data.tip_bin.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d06bb64",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_data.euclidean.hist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10aae180",
   "metadata": {},
   "source": [
    "## 3. Generate raw data schema\n",
    "\n",
    "The [TensorFlow Data Validation (TFDV)](https://www.tensorflow.org/tfx/data_validation/get_started) data schema will be used in:\n",
    "1. Identify the raw data types and shapes in the data transformation.\n",
    "2. Create the serving input signature for the custom model.\n",
    "3. Validate the new raw training data in the TFX pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4087d5fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = tfdv.generate_statistics_from_dataframe(\n",
    "    dataframe=sample_data,\n",
    "    stats_options=tfdv.StatsOptions(\n",
    "        label_feature='tip_bin',\n",
    "        weight_feature=None,\n",
    "        sample_rate=1,\n",
    "        num_top_values=50\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "091fbd77",
   "metadata": {},
   "outputs": [],
   "source": [
    "tfdv.visualize_statistics(stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d251e09b",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = tfdv.infer_schema(statistics=stats)\n",
    "tfdv.display_schema(schema=schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "502c49f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_schema_location = os.path.join(RAW_SCHEMA_DIR, 'schema.pbtxt')\n",
    "tfdv.write_schema_text(schema, raw_schema_location)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59df0723",
   "metadata": {},
   "source": [
    "## 4. Create Vertex Dataset resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d9b605",
   "metadata": {},
   "outputs": [],
   "source": [
    "vertex_ai.init(\n",
    "    project=PROJECT,\n",
    "    location=REGION\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "464ab920",
   "metadata": {},
   "source": [
    "### Create the dataset resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a1707a",
   "metadata": {},
   "outputs": [],
   "source": [
    "bq_uri = f\"bq://{PROJECT}.{BQ_DATASET_NAME}.{BQ_TABLE_NAME}\"\n",
    "\n",
    "dataset = vertex_ai.TabularDataset.create(\n",
    "    display_name=DATASET_DISPLAY_NAME, bq_source=bq_uri)\n",
    "\n",
    "dataset.gca_resource"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c761fadb",
   "metadata": {},
   "source": [
    "### Get the dataset resource\n",
    "\n",
    "The dataset resource is retrieved by display name. Because multiple datasets can have the same display name, we retrieve the most recent updated one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d78b7f4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = vertex_ai.TabularDataset.list(\n",
    "    filter=f\"display_name={DATASET_DISPLAY_NAME}\", \n",
    "    order_by=\"update_time\")[-1]\n",
    "\n",
    "print(\"Dataset resource name:\", dataset.resource_name)\n",
    "print(\"Dataset BigQuery source:\", dataset.gca_resource.metadata['inputConfig']['bigquerySource']['uri'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "569249de",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "common-cpu.m79",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m79"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}