notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8EcdxqUnftBL"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "rbTKH4_6f1ux"
      },
      "source": [
        "## Online feature serving and vector retrieval of BigQuery data with Vertex AI Feature Store\n",
        "\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.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://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.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",
        "  <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/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WWNCLbZZ6MLi"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use `Vertex AI Feature Store` for online serving and vector retrieval of feature values in `BigQuery`.\n",
        "\n",
        "Learn more about [Vertex AI Feature Store](https://cloud.google.com/vertex-ai/docs/featurestore/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CBeo3dIqJVDd"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to create and use an online feature store instance to host and serve data in `BigQuery` with `Vertex AI Feature Store` in an end to end workflow of features serving and vector retrieval user journey.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- `Vertex AI Feature Store`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Provision an online feature store instance to host and serve data.\n",
        "- Create an online feature store instance to serve a `BigQuery` table.\n",
        "- Use the online server to search nearest neighbors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H1t0xA5XB_3n"
      },
      "source": [
        "###Note\n",
        "This is a public Preview release. By using the feature, you acknowledge that you're aware of the open issues and that this preview is provided “as is” under the pre-GA terms of service.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbCa7Pcpqgaz"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial uses the [Google Patents Public Data](https://pantheon.corp.google.com/marketplace/product/google_patents_public_datasets/google-patents-public-data) dataset from the `BigQuery` public datasets.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nrggkyCUrhZM"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* `Vertex AI`\n",
        "* `BigQuery`\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing)\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTJiDCrYsOmT"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GAqbE5Z2sTVM"
      },
      "outputs": [],
      "source": [
        "# Install the packages\n",
        "! pip3 install --upgrade --quiet google-cloud-aiplatform\\\n",
        "                                 google-cloud-bigquery\\\n",
        "                                 google-cloud-bigquery-datatransfer\\\n",
        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "np60_uuCs7X5"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_u0aEgaSs-3v"
      },
      "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": "mdqw6ADTtJRI"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bFANidV0tPbo"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t5cutPRQtQ7m"
      },
      "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": "lfY8yWnbtZ0K"
      },
      "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": "G6iMMALZthFM"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ni5jx6RGtzG3"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "deb54cbe532e"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your service account]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UzsMphY2t4-v"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R_OnJm_Yt8bw"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x4ybIfzhuAOc"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gmnRqX6BuBnx"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uN9JoC1buE9P"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YNAvMVJjuH5b"
      },
      "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": "code",
      "execution_count": null,
      "metadata": {
        "id": "90a7c39da629"
      },
      "outputs": [],
      "source": [
        "! gcloud auth list"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OMCl0avIusKl"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifbIQuN1uz2r"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform, bigquery\n",
        "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
        "from vertexai.resources.preview import FeatureOnlineStore, FeatureView\n",
        "from vertexai.resources.preview.feature_store import utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ooJNSOvu6Q5"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mLTm3pquu9ar"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "\n",
        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "szHWlytYTA0I"
      },
      "source": [
        "## Set up data source in BigQuery\n",
        "\n",
        "### Requirements\n",
        "The data source has to be a BigQuery table or a BigQuery view, with the following requirements on columns:\n",
        "1. [*Required*] One entity id column, type: string\n",
        "2. [*Required*] One embedding column, type: double array\n",
        "3. [*Optional*] One or more filtering columns, type: string or string array\n",
        "4. [*Optional*] One crowding column, type: integer. Crowding ensures that results are diverse by returning at most k' < k neighbors with any single crowding attribute out of k total neighbors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E2i4ZmnrT1lf"
      },
      "source": [
        "### Test data source\n",
        "\n",
        "Select a subset and exclude the repeated records type of columns that are not compatible with Feature Store from the `patents-public-data.google_patents_research.publications_202304` table:\n",
        "\n",
        "Create a small dataset (<=100MB) for demo purposes, you can use the full dataset if needed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IYD5pBRL2fvX"
      },
      "outputs": [],
      "source": [
        "FEATURE_EXTRACT_QUERY_FULL = \"\"\"\n",
        "SELECT publication_number, embedding_v1 as embedding, url, country, publication_description,\n",
        "cpc_low, cpc_inventive_low, top_terms, title, CAST(title_translated as INT) as title_translated,\n",
        "abstract, CAST(abstract_translated as INT) as abstract_translated,\n",
        "cited_by[safe_offset(0)].filing_date as cited_by_filing_date,\n",
        "similar[safe_offset(0)].filing_date as similar_filing_date\n",
        "FROM `patents-public-data.google_patents_research.publications_202304`\n",
        "\"\"\"\n",
        "FEATURE_EXTRACT_QUERY_SMALL = f\"{FEATURE_EXTRACT_QUERY_FULL} WHERE cited_by[safe_offset(0)].filing_date is not NULL LIMIT 1000\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZVec1FjhpwP6"
      },
      "source": [
        "This data source has some filtering columns (e.g. country) and crowding columns (e.g. cited_by_filing_date). Below is the schema of publications data table we use in this guide:\n",
        "\n",
        "|Column name |  Type  |   Mode   |\n",
        "|------------|--------|----------|\n",
        "|publication_number\t| STRING |\tNULLABLE\n",
        "|embedding |\tFLOAT |\tREPEATED\n",
        "|url |\tSTRING\t| NULLABLE\n",
        "|country|\tSTRING|\tNULLABLE|\n",
        "|publication_description|\tSTRING|\tNULLABLE|\n",
        "|cpc_low|\tSTRING|\tREPEATED|\n",
        "|cpc_inventive_low|\tSTRING|\tREPEATED|\n",
        "|top_terms|\tSTRING|\tREPEATED|\n",
        "|title|\tSTRING|\tNULLABLE|\n",
        "|title_translated|\tINTEGER|\tNULLABLE|\n",
        "|abstract|\tSTRING|\tNULLABLE|\n",
        "|abstract_translated|\tINTEGER|\tNULLABLE|\n",
        "|cited_by_filing_date|\tINTEGER|\tNULLABLE|\n",
        "|similar_filing_date|\tINTEGER|\tNULLABLE|"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZIsNo40p2Ns"
      },
      "source": [
        "View the retrieved data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qhp2A2wCpwrQ"
      },
      "outputs": [],
      "source": [
        "bq_client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "product_data = bq_client.query(FEATURE_EXTRACT_QUERY_SMALL).result().to_dataframe()\n",
        "\n",
        "print(product_data.shape)\n",
        "product_data.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7CspnYk2GyDV"
      },
      "source": [
        "### Create BigQuery dataset\n",
        "\n",
        "Create a BigQuery dataset to hold the BigQuery table for the tutorial. The dataset must be in the same region as your feature online store instance. Since the source data for this tutorial is located in the `US` region, you can copy the dataset to the desired region. If you use your own data and dataset, you can also use that dataset to create the BigQuery table.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0-2pb-NUGzMI"
      },
      "outputs": [],
      "source": [
        "def create_bq_dataset(datasetId, region):\n",
        "    dataset = bigquery.Dataset(f\"{PROJECT_ID}.{datasetId}\")\n",
        "    dataset.location = region\n",
        "    dataset = bq_client.create_dataset(\n",
        "        dataset, exists_ok=True, timeout=30\n",
        "    )  # Make an API request.\n",
        "\n",
        "    # Confirm dataset created.\n",
        "    print(f\"Created dataset {dataset} in region {region}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "41b9122c8fda"
      },
      "outputs": [],
      "source": [
        "# First, create a dataset to keep the feature store source data if it does not already exist.\n",
        "BQ_DATASET_ID = (\n",
        "    f\"featurestore_demo_{REGION.replace('-', '_')}\"  # @param {type:\"string\"}\n",
        ")\n",
        "create_bq_dataset(BQ_DATASET_ID, REGION)\n",
        "\n",
        "# The source data for this demo is located in the US region, so a temp dataset must also be located in the US region to transfer data to the desired region.\n",
        "TIMESTAMP = !date +%s\n",
        "BQ_TEMP_DATASET_ID = f\"fs_temp_{TIMESTAMP[0]}\"\n",
        "create_bq_dataset(BQ_TEMP_DATASET_ID, \"US\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1KSrpN_LG1ia"
      },
      "source": [
        "#### Create a BigQuery table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JXP2ZaFAG18W"
      },
      "outputs": [],
      "source": [
        "# 1. execute the query and keep results in temp dataset\n",
        "BQ_TABLE_ID = \"publications_202304_small\"  # @param {type:\"string\"}\n",
        "BQ_TEMP_TABLE_ID_FQN = f\"{PROJECT_ID}.{BQ_TEMP_DATASET_ID}.{BQ_TABLE_ID}\"\n",
        "job_config = bigquery.QueryJobConfig(destination=BQ_TEMP_TABLE_ID_FQN)\n",
        "query_job = bq_client.query(FEATURE_EXTRACT_QUERY_SMALL, job_config=job_config)\n",
        "\n",
        "try:\n",
        "    query_job.result()\n",
        "except Exception as e:\n",
        "    # Table already exists\n",
        "    print(\"Error: \", e.message)\n",
        "print(f\"Created table: {BQ_TEMP_TABLE_ID_FQN}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bab92f3c19e8"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import time\n",
        "\n",
        "from google.cloud import bigquery_datatransfer\n",
        "\n",
        "# 2. copy temp dataset to the desired region\n",
        "transfer_client = bigquery_datatransfer.DataTransferServiceClient()\n",
        "\n",
        "transfer_config = bigquery_datatransfer.TransferConfig(\n",
        "    destination_dataset_id=BQ_DATASET_ID,\n",
        "    display_name=f\"Copy BQ {BQ_TEMP_DATASET_ID} to {BQ_DATASET_ID}\",\n",
        "    data_source_id=\"cross_region_copy\",\n",
        "    params={\n",
        "        \"source_project_id\": PROJECT_ID,\n",
        "        \"source_dataset_id\": BQ_TEMP_DATASET_ID,\n",
        "    },\n",
        "    schedule=\"None\",\n",
        "    schedule_options=bigquery_datatransfer.ScheduleOptions(disable_auto_scheduling=True),\n",
        ")\n",
        "transfer_config = transfer_client.create_transfer_config(\n",
        "    parent=transfer_client.common_project_path(PROJECT_ID),\n",
        "    transfer_config=transfer_config,\n",
        ")\n",
        "print(\n",
        "    f\"Created data transfer config for dataset copy: {transfer_config.name}\"\n",
        ")\n",
        "\n",
        "transfer_run = transfer_client.start_manual_transfer_runs(request=bigquery_datatransfer.StartManualTransferRunsRequest(\n",
        "    parent=transfer_config.name,\n",
        "    requested_run_time=datetime.datetime.now(tz=datetime.timezone.utc)\n",
        ")).runs[0]\n",
        "\n",
        "print(\n",
        "    f\"Started running dataset copy job: {transfer_run.name}\"\n",
        ")\n",
        "\n",
        "job_completed = False\n",
        "while True:\n",
        "    run = transfer_client.get_transfer_run(bigquery_datatransfer.GetTransferRunRequest(name=transfer_run.name))\n",
        "    if run.state == bigquery_datatransfer.TransferState.SUCCEEDED:\n",
        "        job_completed = True\n",
        "        print(\"Dataset copy is complete.\")\n",
        "        # clean up temp dataset.\n",
        "        transfer_client.delete_transfer_config(\n",
        "            request=bigquery_datatransfer.DeleteTransferConfigRequest(\n",
        "                name=transfer_config.name\n",
        "            )\n",
        "        )\n",
        "        bq_client.delete_dataset(\n",
        "            bigquery.Dataset(f\"{PROJECT_ID}.{BQ_TEMP_DATASET_ID}\"),\n",
        "            delete_contents=True,\n",
        "        )\n",
        "        # confirm table is created.\n",
        "        BQ_TABLE_ID_FQN = f\"{PROJECT_ID}.{BQ_DATASET_ID}.{BQ_TABLE_ID}\"\n",
        "        table = bq_client.get_table(bigquery.Table(BQ_TABLE_ID_FQN))\n",
        "        DATA_SOURCE = f\"bq://{table}\"\n",
        "        print(f\"Data source is: {DATA_SOURCE}\")\n",
        "    elif run.state == bigquery_datatransfer.TransferState.FAILED:\n",
        "        job_completed = True\n",
        "        print(f\"Dataset copy is failed. \\n {run}\")\n",
        "    if job_completed is True:\n",
        "        break\n",
        "    time.sleep(30)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AYzwb095DJTl"
      },
      "source": [
        "## Set up and start online serving\n",
        "\n",
        "To serve embedding data in feature store, you need to do the following:\n",
        "\n",
        "1. Create an online store cluster to host the data.\n",
        "2. Define the data (FeatureView) to be served by the newly-created instance."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "43CSdWFTTn-_"
      },
      "source": [
        "### Create Feature Online Store\n",
        "\n",
        "Create an Optimized feature online store. After the long-running operation (LRO) is complete, the result will be shown.\n",
        "\n",
        "> **Note:** This operation might take up to 10 minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e544c2eb906f"
      },
      "outputs": [],
      "source": [
        "FEATURE_ONLINE_STORE_ID = \"my_feature_online_store_unique\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uoCExlfzfjNE"
      },
      "outputs": [],
      "source": [
        "my_fos = FeatureOnlineStore.create_optimized_store(FEATURE_ONLINE_STORE_ID)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7240bcafef77"
      },
      "outputs": [],
      "source": [
        "# get full information of the created feature online store instance\n",
        "my_fos.gca_resource"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ldCMH-IOOx-G"
      },
      "source": [
        "#### Verify that the `FeatureOnlineStore` instance is created by getting the online stores instance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D3dZEtXKIqXT"
      },
      "outputs": [],
      "source": [
        "# Use get to verify the store is created.\n",
        "FeatureOnlineStore(FEATURE_ONLINE_STORE_ID).gca_resource"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2vIjCzCwIGMF"
      },
      "source": [
        "#### List all online stores for the location"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rk8JseF4ICur"
      },
      "outputs": [],
      "source": [
        "# Use list to verify the store is created.\n",
        "all_foss = FeatureOnlineStore.list()\n",
        "for fos in all_foss:\n",
        "    print(fos.gca_resource)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf0f5b6c204c"
      },
      "source": [
        "#### Optional: Use an online store created in the location"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3d01ba11e964"
      },
      "outputs": [],
      "source": [
        "# another_fos = FeatureOnlineStore(\"an_existing_feature_online_store\")  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7QZnKMo41ud5"
      },
      "source": [
        "### Create feature view instance\n",
        "\n",
        "After creating a `FeatureOnlineStore` instance, you define the features to serve with it. To do this, create a `FeatureView` instance, which specifies the following:\n",
        "\n",
        "* A data source (BigQuery table or view URI or FeatureGroup/features ) synced to the `FeatureOnlineStore` instance for serving.\n",
        "* The cron schedule to run the sync pipeline.\n",
        "\n",
        "Within feature view creation, a sync job will be scheduled, either started immediately or following the cron schedule. In the sync job, data is exported to Cloud Bigtable, index is built and deployed to GKE cluster."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9EfVmeu5KBu"
      },
      "outputs": [],
      "source": [
        "FEATURE_VIEW_ID = \"feature_view_publications\"  # @param {type: \"string\"}\n",
        "# A schedule will be created based on cron setting.\n",
        "# If cron is empty, an immediate schedule job will be started.\n",
        "CRON_SCHEDULE = \"TZ=America/Los_Angeles 00 13 11 8 *\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8VQ4P2om6SNQ"
      },
      "outputs": [],
      "source": [
        "# Index configs\n",
        "DIMENSIONS = 64  # @param {type: \"number\"}\n",
        "EMBEDDING_COLUMN = \"embedding\"  # @param {type: \"string\"}\n",
        "# Optional\n",
        "LEAF_NODE_EMBEDDING_COUNT = 10000  # @param {type: \"number\"}\n",
        "# Optional\n",
        "CROWDING_COLUMN = \"cited_by_filing_date\"  # @param {type: \"string\"}\n",
        "# Optional\n",
        "FILTER_COLUMNS = [\"country\"]  # @param"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AKdSax3N1ypm"
      },
      "outputs": [],
      "source": [
        "big_query_source = FeatureView.BigQuerySource(\n",
        "    uri=DATA_SOURCE, entity_id_columns=[\"publication_number\"]\n",
        ")\n",
        "\n",
        "index_config = utils.IndexConfig(\n",
        "    embedding_column=EMBEDDING_COLUMN,\n",
        "    crowding_column=CROWDING_COLUMN,\n",
        "    filter_column=FILTER_COLUMNS,\n",
        "    dimentions=DIMENSIONS,\n",
        "    distance_measure_type=utils.DistanceMeasureType.DOT_PRODUCT_DISTANCE,\n",
        "    algorithm_config=utils.TreeAhConfig(),\n",
        ")\n",
        "\n",
        "print(f\"index_config: {index_config}\")\n",
        "\n",
        "my_fv = my_fos.create_feature_view_from_big_query(\n",
        "    FEATURE_VIEW_ID,\n",
        "    big_query_source=big_query_source,\n",
        "    sync_config=CRON_SCHEDULE,  # Optional, can be set to None.\n",
        "    index_config=index_config,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a10cd8e6061a"
      },
      "outputs": [],
      "source": [
        "# get full information of the created feature view\n",
        "my_fv.gca_resource"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mmYQZroEO7dp"
      },
      "source": [
        "### Verify feature view instance creation\n",
        "\n",
        "Verify that the FeatureView instance is created by gettting the feature view."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PaqTIYuwPFC7"
      },
      "outputs": [],
      "source": [
        "FeatureView(\n",
        "    FEATURE_VIEW_ID, feature_online_store_id=FEATURE_ONLINE_STORE_ID\n",
        ").gca_resource"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sqpu4nHAO4pW"
      },
      "source": [
        "Verify that the FeatureView instance is created by listing all the feature views within the online store."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O_tv0Z3BPHW4"
      },
      "outputs": [],
      "source": [
        "all_fvs = FeatureView.list(feature_online_store_id=FEATURE_ONLINE_STORE_ID)\n",
        "for fv in all_fvs:\n",
        "    print(fv.gca_resource)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KZglLt0iuVs0"
      },
      "source": [
        "### Feature view syncs\n",
        "\n",
        "The sync pipeline executes according to the schedule specified in the `FeatureView` instance.\n",
        "\n",
        "To skip the wait and execute the sync pipeline immediately, start the sync manually."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G-j5B6InuNfW"
      },
      "outputs": [],
      "source": [
        "sync_response = my_fv.sync()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GqyYqYplMUjZ"
      },
      "source": [
        "The `sync_response` contains the ID of the sync job.\n",
        "\n",
        "#### Use `get_feature_view_sync` to check the status of the job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_kl8ZBAR2o0b"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "while True:\n",
        "    feature_view_sync = my_fv.get_sync(\n",
        "        sync_response.resource_name.split(\"/\")[9]\n",
        "    ).gca_resource\n",
        "    if feature_view_sync.run_time.end_time.seconds > 0:\n",
        "        status = \"Succeed\" if feature_view_sync.final_status.code == 0 else \"Failed\"\n",
        "        print(f\"Sync {status} for {feature_view_sync.name}. \\n {feature_view_sync}\")\n",
        "        # wait a little more for the job to properly shutdown\n",
        "        time.sleep(30)\n",
        "        break\n",
        "    else:\n",
        "        print(\"Sync ongoing, waiting for 30 seconds.\")\n",
        "    time.sleep(30)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3XHiHnAIJGj"
      },
      "source": [
        "#### Use `list_feature_view_syncs` to view all your syncs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3CU67MyLIOAk"
      },
      "outputs": [],
      "source": [
        "fv_syncs = fv.list_syncs()\n",
        "for fv_sync in fv_syncs:\n",
        "    print(f\" * {fv_sync.gca_resource}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eUBuFxXBbYOV"
      },
      "source": [
        "### Start online serving\n",
        "\n",
        "After the data sync is complete, use the `FetchFeatureValuesRequest` and `SearchNearestEntities` APIs to retrieve the data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jd7gHqZ8blJ-"
      },
      "source": [
        "#### Set `NearestNeighborQuery.StringFilter`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4erlJr3_1Du"
      },
      "outputs": [],
      "source": [
        "results_df = (\n",
        "    bq_client.query(f\"select publication_number from {BQ_TABLE_ID_FQN} limit 1\")\n",
        "    .result()\n",
        "    .to_dataframe()\n",
        ")\n",
        "ENTITY_ID = results_df.loc[0, \"publication_number\"]\n",
        "print(f\"Sample publication number: {ENTITY_ID}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7sYPMibU3YwR"
      },
      "outputs": [],
      "source": [
        "country_filter = NearestNeighborQuery.StringFilter(\n",
        "    name=\"country\",\n",
        "    allow_tokens=[\"WIPO (PCT)\"],  # try different allow tokens\n",
        "    deny_tokens=[\"United States\"],  # try different deny tokens\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v5gYpORiBW1n"
      },
      "source": [
        "#### Search with `ENTITY_ID`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NWhrk3pqSU6h"
      },
      "outputs": [],
      "source": [
        "my_fv.search(\n",
        "    entity_id=ENTITY_ID,\n",
        "    neighbor_count=5,\n",
        "    string_filters=[country_filter],\n",
        "    return_full_entity=True,  # returning entities with metadata\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OQEpt08GBX-b"
      },
      "source": [
        "#### Search with `Embedding`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JWcOT26krxuj"
      },
      "outputs": [],
      "source": [
        "EMBEDDINGS = [1] * DIMENSIONS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e3DZl0xRAusK"
      },
      "outputs": [],
      "source": [
        "my_fv.search(\n",
        "    embedding_value=EMBEDDINGS,\n",
        "    neighbor_count=10,\n",
        "    string_filters=[country_filter],\n",
        "    return_full_entity=True,  # returning entities with metadata\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MKALOxbsZfce"
      },
      "source": [
        "#### Use the `FetchFeatureValues` API to retrieve the full data without search\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_7op0efLdDFs"
      },
      "outputs": [],
      "source": [
        "my_fv.read(key=[ENTITY_ID])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OR1ve2OyyiKq"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Um27J_vvvzGc"
      },
      "outputs": [],
      "source": [
        "# Delete Feature View\n",
        "my_fv.delete()\n",
        "\n",
        "# Delete Feature Online Store\n",
        "my_fos.delete(force=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}