notebooks/official/feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "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 style=\"text-align: center\">\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=\"Google Colaboratory 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%2Ffeature_store%2Fonline_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.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/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\"><br> Open in 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/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\"><br> View on GitHub\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://console.cloud.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": "7cc240e9c1a6"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTJiDCrYsOmT"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "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",
        "                                 db-dtypes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "np60_uuCs7X5"
      },
      "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": "_u0aEgaSs-3v"
      },
      "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": "e669f8088ac3"
      },
      "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": "60763ee24ce0"
      },
      "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": "t5cutPRQtQ7m"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type: \"string\"}\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
        "\n",
        "API_ENDPOINT = f\"{LOCATION}-aiplatform.googleapis.com\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OMCl0avIusKl"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifbIQuN1uz2r"
      },
      "outputs": [],
      "source": [
        "from google.cloud import bigquery\n",
        "from google.cloud.aiplatform_v1.types import NearestNeighborQuery\n",
        "from vertexai.resources.preview import (FeatureOnlineStore, FeatureView,\n",
        "                                        FeatureViewBigQuerySource)\n",
        "from vertexai.resources.preview.feature_store import utils"
      ]
    },
    {
      "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 (for example, 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 doesn't already exist.\n",
        "BQ_DATASET_ID = \"featurestore_demo_us\"  # @param {type:\"string\"}\n",
        "create_bq_dataset(BQ_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": [
        "# Second, execute the query and store the results into a table\n",
        "BQ_TABLE_ID = \"publications_202304_small\"  # @param {type:\"string\"}\n",
        "BQ_TABLE_ID_FQN = f\"{PROJECT_ID}.{BQ_DATASET_ID}.{BQ_TABLE_ID}\"\n",
        "\n",
        "job_config = bigquery.QueryJobConfig(destination=BQ_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",
        "\n",
        "print(f\"Created table: {BQ_TABLE_ID}\")\n",
        "DATA_SOURCE = f\"bq://{BQ_TABLE_ID_FQN}\""
      ]
    },
    {
      "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 is 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 is scheduled, either started immediately or following the cron schedule. In the sync job, data is exported, index is built and deployed to Feature Store backend."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9EfVmeu5KBu"
      },
      "outputs": [],
      "source": [
        "FEATURE_VIEW_ID = \"feature_view_publications\"  # @param {type: \"string\"}\n",
        "# A schedule is created based on cron setting.\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 = FeatureViewBigQuerySource(\n",
        "    uri=DATA_SOURCE, entity_id_columns=[\"publication_number\"]\n",
        ")\n",
        "\n",
        "index_config = utils.IndexConfig(\n",
        "    embedding_column=EMBEDDING_COLUMN,\n",
        "    dimensions=DIMENSIONS,\n",
        "    crowding_column=CROWDING_COLUMN,\n",
        "    filter_columns=FILTER_COLUMNS,\n",
        "    algorithm_config=utils.TreeAhConfig(),\n",
        ")\n",
        "\n",
        "print(f\"index_config: {index_config}\")\n",
        "\n",
        "my_fv = my_fos.create_feature_view(\n",
        "    FEATURE_VIEW_ID,\n",
        "    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 getting 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": [
        "# It takes some time for the DNS to be fully ready\n",
        "time.sleep(300)\n",
        "\n",
        "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)\n",
        "\n",
        "# Delete BigQuery dataset. Uncomment and run the command below if you want to delete the BigQuery set.\n",
        "# Do this only if the dataset is created for this demo.\n",
        "# bq_client.delete_dataset(BQ_DATASET_ID, delete_contents=True, not_found_ok=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
}