Feast on Kubeflow Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Minimal Ride Hailing Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Restart your Kernel after installing these packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install protobuf gcsfs feast -U -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating gs://feast-staging-bucket-4179836/...\r\n"
     ]
    }
   ],
   "source": [
    "import random \n",
    "staging_bucket = f'gs://feast-staging-bucket-{random.randint(1000000, 10000000)}/'\n",
    "!gsutil mb {staging_bucket}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Imports and Feast Client initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from feast import Client, Feature, Entity, ValueType, FeatureTable\n",
    "from feast.data_source import FileSource, KafkaSource\n",
    "from feast.data_format import ParquetFormat, AvroFormat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(\n",
    "    core_url=\"kf-feast-core.kubeflow.svc:6565\",\n",
    "    serving_url=\"feast-online-serving.kubeflow.svc:6566\",\n",
    "    spark_launcher=\"k8s\",\n",
    "    spark_staging_location=staging_bucket,\n",
    "    spark_k8s_namespace=\"kubeflow\",\n",
    "    redis_host=\"kf-feast-redis-headless.kubeflow.svc\",\n",
    "    historical_feature_output_location=f\"{staging_bucket}historical\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Declare Features and Entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver_id = Entity(name=\"driver_id\", description=\"Driver identifier\", value_type=ValueType.INT64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Daily updated features \n",
    "acc_rate = Feature(\"acc_rate\", ValueType.FLOAT)\n",
    "conv_rate = Feature(\"conv_rate\", ValueType.FLOAT)\n",
    "avg_daily_trips = Feature(\"avg_daily_trips\", ValueType.INT32)\n",
    "\n",
    "# Real-time updated features\n",
    "trips_today = Feature(\"trips_today\", ValueType.INT32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://feast-staging-bucket-4179836/test_data\n"
     ]
    }
   ],
   "source": [
    "# Offline data will be stored in this location\n",
    "demo_data_location = os.path.join(os.getenv(\"FEAST_SPARK_STAGING_LOCATION\", staging_bucket), \"test_data\")\n",
    "print(demo_data_location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver_statistics_source_uri = os.path.join(demo_data_location, \"driver_statistics\")\n",
    "\n",
    "driver_statistics = FeatureTable(\n",
    "    name = \"driver_statistics\",\n",
    "    entities = [\"driver_id\"],\n",
    "    features = [\n",
    "        acc_rate,\n",
    "        conv_rate,\n",
    "        avg_daily_trips\n",
    "    ],\n",
    "    batch_source=FileSource(\n",
    "        event_timestamp_column=\"datetime\",\n",
    "        created_timestamp_column=\"created\",\n",
    "        file_format=ParquetFormat(),\n",
    "        file_url=driver_statistics_source_uri,\n",
    "        date_partition_column=\"date\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver_trips_source_uri = os.path.join(demo_data_location, \"driver_trips\")\n",
    "\n",
    "\n",
    "driver_trips = FeatureTable(\n",
    "    name = \"driver_trips\",\n",
    "    entities = [\"driver_id\"],\n",
    "    features = [\n",
    "        trips_today\n",
    "    ],\n",
    "    batch_source=FileSource(\n",
    "        event_timestamp_column=\"datetime\",\n",
    "        created_timestamp_column=\"created\",\n",
    "        file_format=ParquetFormat(),\n",
    "        file_url=driver_trips_source_uri,\n",
    "        date_partition_column=\"date\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Registering entities and feature tables in Feast Core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.apply(driver_id)\n",
    "client.apply(driver_statistics)\n",
    "client.apply(driver_trips)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "spec:\n",
      "  name: driver_statistics\n",
      "  entities:\n",
      "  - driver_id\n",
      "  features:\n",
      "  - name: avg_daily_trips\n",
      "    valueType: INT32\n",
      "  - name: conv_rate\n",
      "    valueType: FLOAT\n",
      "  - name: acc_rate\n",
      "    valueType: FLOAT\n",
      "  batchSource:\n",
      "    type: BATCH_FILE\n",
      "    eventTimestampColumn: datetime\n",
      "    datePartitionColumn: date\n",
      "    createdTimestampColumn: created\n",
      "    fileOptions:\n",
      "      fileFormat:\n",
      "        parquetFormat: {}\n",
      "      fileUrl: gs://feast-staging-bucket-4179836/test_data/driver_statistics\n",
      "meta:\n",
      "  createdTimestamp: '2021-02-03T05:50:53Z'\n",
      "\n",
      "spec:\n",
      "  name: driver_trips\n",
      "  entities:\n",
      "  - driver_id\n",
      "  features:\n",
      "  - name: trips_today\n",
      "    valueType: INT32\n",
      "  batchSource:\n",
      "    type: BATCH_FILE\n",
      "    eventTimestampColumn: datetime\n",
      "    datePartitionColumn: date\n",
      "    createdTimestampColumn: created\n",
      "    fileOptions:\n",
      "      fileFormat:\n",
      "        parquetFormat: {}\n",
      "      fileUrl: gs://feast-staging-bucket-4179836/test_data/driver_trips\n",
      "meta:\n",
      "  createdTimestamp: '2021-02-03T05:50:53Z'\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(client.get_feature_table(\"driver_statistics\").to_yaml())\n",
    "print(client.get_feature_table(\"driver_trips\").to_yaml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Populating batch source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feast is agnostic to how the batch source is populated, as long as it complies to the Feature Table specification. Therefore, any existing ETL tools can be used for the purpose of data ingestion. Alternatively, you can also use Feast SDK to ingest a Panda Dataframe to the batch source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_entities():\n",
    "    return np.random.choice(999999, size=100, replace=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_trips(entities):\n",
    "    df = pd.DataFrame(columns=[\"driver_id\", \"trips_today\", \"datetime\", \"created\"])\n",
    "    df['driver_id'] = entities\n",
    "    df['trips_today'] = np.random.randint(0, 1000, size=100).astype(np.int32)\n",
    "    df['datetime'] = pd.to_datetime(\n",
    "            np.random.randint(\n",
    "                datetime(2020, 10, 10).timestamp(),\n",
    "                datetime(2020, 10, 20).timestamp(),\n",
    "                size=100),\n",
    "        unit=\"s\"\n",
    "    )\n",
    "    df['created'] = pd.to_datetime(datetime.now())\n",
    "    return df\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_stats(entities):\n",
    "    df = pd.DataFrame(columns=[\"driver_id\", \"conv_rate\", \"acc_rate\", \"avg_daily_trips\", \"datetime\", \"created\"])\n",
    "    df['driver_id'] = entities\n",
    "    df['conv_rate'] = np.random.random(size=100).astype(np.float32)\n",
    "    df['acc_rate'] = np.random.random(size=100).astype(np.float32)\n",
    "    df['avg_daily_trips'] = np.random.randint(0, 1000, size=100).astype(np.int32)\n",
    "    df['datetime'] = pd.to_datetime(\n",
    "            np.random.randint(\n",
    "                datetime(2020, 10, 10).timestamp(),\n",
    "                datetime(2020, 10, 20).timestamp(),\n",
    "                size=100),\n",
    "        unit=\"s\"\n",
    "    )\n",
    "    df['created'] = pd.to_datetime(datetime.now())\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "entities = generate_entities()\n",
    "stats_df = generate_stats(entities)\n",
    "trips_df = generate_trips(entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing temporary file(s)...\n",
      "Data has been successfully ingested into FeatureTable batch source.\n",
      "Removing temporary file(s)...\n",
      "Data has been successfully ingested into FeatureTable batch source.\n"
     ]
    }
   ],
   "source": [
    "client.ingest(driver_statistics, stats_df)\n",
    "client.ingest(driver_trips, trips_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Historical Retrieval For Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a training dataset from offline feature tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gcsfs\n",
    "from pyarrow.parquet import ParquetDataset\n",
    "from urllib.parse import urlparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>driver_id</th>\n",
       "      <th>event_timestamp</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>441896</td>\n",
       "      <td>2020-10-18 06:02:29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>59256</td>\n",
       "      <td>2020-10-19 06:02:31</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>79041</td>\n",
       "      <td>2020-10-19 00:58:04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>210520</td>\n",
       "      <td>2020-10-18 00:51:04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>552755</td>\n",
       "      <td>2020-10-18 19:15:34</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>991178</td>\n",
       "      <td>2020-10-18 07:27:46</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>577396</td>\n",
       "      <td>2020-10-19 10:46:28</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>401906</td>\n",
       "      <td>2020-10-19 10:58:32</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>135710</td>\n",
       "      <td>2020-10-18 08:22:22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>849882</td>\n",
       "      <td>2020-10-19 07:20:05</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   driver_id     event_timestamp\n",
       "0     441896 2020-10-18 06:02:29\n",
       "1      59256 2020-10-19 06:02:31\n",
       "2      79041 2020-10-19 00:58:04\n",
       "3     210520 2020-10-18 00:51:04\n",
       "4     552755 2020-10-18 19:15:34\n",
       "5     991178 2020-10-18 07:27:46\n",
       "6     577396 2020-10-19 10:46:28\n",
       "7     401906 2020-10-19 10:58:32\n",
       "8     135710 2020-10-18 08:22:22\n",
       "9     849882 2020-10-19 07:20:05"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entities_with_timestamp = pd.DataFrame(columns=['driver_id', 'event_timestamp'])\n",
    "entities_with_timestamp['driver_id'] = np.random.choice(entities, 10, replace=False)\n",
    "entities_with_timestamp['event_timestamp'] = pd.to_datetime(np.random.randint(\n",
    "    datetime(2020, 10, 18).timestamp(),\n",
    "    datetime(2020, 10, 20).timestamp(),\n",
    "    size=10), unit='s')\n",
    "entities_with_timestamp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_historical_features will return immediately once the Spark job has been submitted succesfully.\n",
    "job = client.get_historical_features(\n",
    "    feature_refs=[\n",
    "        \"driver_statistics:avg_daily_trips\",\n",
    "        \"driver_statistics:conv_rate\",\n",
    "        \"driver_statistics:acc_rate\",\n",
    "        \"driver_trips:trips_today\"\n",
    "    ], \n",
    "    entity_source=entities_with_timestamp\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_output_file_uri will block until the Spark job is completed.\n",
    "output_file_uri = job.get_output_file_uri()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>driver_id</th>\n",
       "      <th>event_timestamp</th>\n",
       "      <th>driver_statistics__avg_daily_trips</th>\n",
       "      <th>driver_statistics__conv_rate</th>\n",
       "      <th>driver_statistics__acc_rate</th>\n",
       "      <th>driver_trips__trips_today</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>552755</td>\n",
       "      <td>2020-10-18 19:15:34</td>\n",
       "      <td>787</td>\n",
       "      <td>0.601664</td>\n",
       "      <td>0.444756</td>\n",
       "      <td>617.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>441896</td>\n",
       "      <td>2020-10-18 06:02:29</td>\n",
       "      <td>493</td>\n",
       "      <td>0.348591</td>\n",
       "      <td>0.910286</td>\n",
       "      <td>405.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>849882</td>\n",
       "      <td>2020-10-19 07:20:05</td>\n",
       "      <td>534</td>\n",
       "      <td>0.927652</td>\n",
       "      <td>0.397792</td>\n",
       "      <td>988.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>135710</td>\n",
       "      <td>2020-10-18 08:22:22</td>\n",
       "      <td>688</td>\n",
       "      <td>0.477989</td>\n",
       "      <td>0.363104</td>\n",
       "      <td>54.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>577396</td>\n",
       "      <td>2020-10-19 10:46:28</td>\n",
       "      <td>772</td>\n",
       "      <td>0.059130</td>\n",
       "      <td>0.068620</td>\n",
       "      <td>121.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>79041</td>\n",
       "      <td>2020-10-19 00:58:04</td>\n",
       "      <td>261</td>\n",
       "      <td>0.104633</td>\n",
       "      <td>0.109721</td>\n",
       "      <td>888.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>59256</td>\n",
       "      <td>2020-10-19 06:02:31</td>\n",
       "      <td>194</td>\n",
       "      <td>0.170629</td>\n",
       "      <td>0.715083</td>\n",
       "      <td>731.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>401906</td>\n",
       "      <td>2020-10-19 10:58:32</td>\n",
       "      <td>289</td>\n",
       "      <td>0.658645</td>\n",
       "      <td>0.838213</td>\n",
       "      <td>854.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>210520</td>\n",
       "      <td>2020-10-18 00:51:04</td>\n",
       "      <td>951</td>\n",
       "      <td>0.119565</td>\n",
       "      <td>0.845131</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>991178</td>\n",
       "      <td>2020-10-18 07:27:46</td>\n",
       "      <td>632</td>\n",
       "      <td>0.699143</td>\n",
       "      <td>0.345410</td>\n",
       "      <td>607.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   driver_id     event_timestamp  driver_statistics__avg_daily_trips  \\\n",
       "0     552755 2020-10-18 19:15:34                                 787   \n",
       "1     441896 2020-10-18 06:02:29                                 493   \n",
       "2     849882 2020-10-19 07:20:05                                 534   \n",
       "3     135710 2020-10-18 08:22:22                                 688   \n",
       "4     577396 2020-10-19 10:46:28                                 772   \n",
       "5      79041 2020-10-19 00:58:04                                 261   \n",
       "6      59256 2020-10-19 06:02:31                                 194   \n",
       "7     401906 2020-10-19 10:58:32                                 289   \n",
       "8     210520 2020-10-18 00:51:04                                 951   \n",
       "9     991178 2020-10-18 07:27:46                                 632   \n",
       "\n",
       "   driver_statistics__conv_rate  driver_statistics__acc_rate  \\\n",
       "0                      0.601664                     0.444756   \n",
       "1                      0.348591                     0.910286   \n",
       "2                      0.927652                     0.397792   \n",
       "3                      0.477989                     0.363104   \n",
       "4                      0.059130                     0.068620   \n",
       "5                      0.104633                     0.109721   \n",
       "6                      0.170629                     0.715083   \n",
       "7                      0.658645                     0.838213   \n",
       "8                      0.119565                     0.845131   \n",
       "9                      0.699143                     0.345410   \n",
       "\n",
       "   driver_trips__trips_today  \n",
       "0                      617.0  \n",
       "1                      405.0  \n",
       "2                      988.0  \n",
       "3                       54.0  \n",
       "4                      121.0  \n",
       "5                      888.0  \n",
       "6                      731.0  \n",
       "7                      854.0  \n",
       "8                        NaN  \n",
       "9                      607.0  "
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Retrieve the remote training dataset\n",
    "\n",
    "parsed_uri = urlparse(output_file_uri)\n",
    "fs = gcsfs.GCSFileSystem()\n",
    "files = [\"gs://\" + path for path in fs.glob(output_file_uri + '/part-*')]\n",
    "ds = ParquetDataset(files, filesystem=fs)\n",
    "ds.read().to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The retrieved result can now be used for model training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Populating Online Storage with Batch Ingestion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to populate the online storage, we can use Feast SDK to start a Spark batch job which will extract the features from the batch source, then load the features to an online store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = client.start_offline_to_online_ingestion(\n",
    "    driver_statistics,\n",
    "    datetime(2020, 10, 10),\n",
    "    datetime(2020, 10, 20)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<SparkJobStatus.COMPLETED: 3>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# It will take some time before the Spark Job is completed\n",
    "job.get_status()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the job is completed, the SDK can be used to retrieve the result from the online store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'driver_id': 432527},\n",
       " {'driver_id': 66847},\n",
       " {'driver_id': 648263},\n",
       " {'driver_id': 955983},\n",
       " {'driver_id': 650942},\n",
       " {'driver_id': 638553},\n",
       " {'driver_id': 873512},\n",
       " {'driver_id': 480137},\n",
       " {'driver_id': 811349},\n",
       " {'driver_id': 987187}]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entities_sample = np.random.choice(entities, 10, replace=False)\n",
    "entities_sample = [{\"driver_id\": e} for e in entities_sample]\n",
    "entities_sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'driver_statistics:avg_daily_trips': [535,\n",
       "  916,\n",
       "  643,\n",
       "  142,\n",
       "  317,\n",
       "  266,\n",
       "  552,\n",
       "  949,\n",
       "  93,\n",
       "  694],\n",
       " 'driver_id': [432527,\n",
       "  66847,\n",
       "  648263,\n",
       "  955983,\n",
       "  650942,\n",
       "  638553,\n",
       "  873512,\n",
       "  480137,\n",
       "  811349,\n",
       "  987187]}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features = client.get_online_features(\n",
    "    feature_refs=[\"driver_statistics:avg_daily_trips\"],\n",
    "    entity_rows=entities_sample).to_dict()\n",
    "features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>driver_statistics:avg_daily_trips</th>\n",
       "      <th>driver_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>535</td>\n",
       "      <td>432527</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>916</td>\n",
       "      <td>66847</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>643</td>\n",
       "      <td>648263</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>142</td>\n",
       "      <td>955983</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>317</td>\n",
       "      <td>650942</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>266</td>\n",
       "      <td>638553</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>552</td>\n",
       "      <td>873512</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>949</td>\n",
       "      <td>480137</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>93</td>\n",
       "      <td>811349</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>694</td>\n",
       "      <td>987187</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   driver_statistics:avg_daily_trips  driver_id\n",
       "0                                535     432527\n",
       "1                                916      66847\n",
       "2                                643     648263\n",
       "3                                142     955983\n",
       "4                                317     650942\n",
       "5                                266     638553\n",
       "6                                552     873512\n",
       "7                                949     480137\n",
       "8                                 93     811349\n",
       "9                                694     987187"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}