Online Feature Store example notebook

This notebook illustrates the use of Databricks Feature Store to publish features to Databricks Online Tables for real-time serving and automated feature lookup. The problem is to predict the wine quality using a ML model with a variety of static wine features and a realtime input.

This notebook creates an endpoint to predict the quality of a bottle of wine, given an ID and the realtime feature alcohol by volume (ABV).

The notebook is structured as follows:

1. Prepare the feature table.
2. Set up Databricks Online Table.
   • This notebook uses Databricks Online Tables. For a list of supported functionality, see the Databricks documentation (AWS | Azure).
3. Train and deploy the model.
4. Serve realtime queries with automatic feature lookup.
5. Clean up.

Data Set

This example uses the Wine Quality Data Set.

Requirements

• Serverless Compute (AWS | Azure) (recommended)
• Classic compute running Databricks Runtime 14.2 for Machine Learning or above.

Prepare the feature table

Suppose you need to build an endpoint to predict wine quality with just the wine_id. This requires a feature table saved in Feature Store where the endpoint can look up features of the wine by the wine_id. For the purpose of this demo, we need to prepare this feature table ourselves first. The steps are:

1. Load and clean the raw data.
2. Separate features and labels.
3. Save features into a feature table.

Load and clean the raw data

The raw data contains 12 columns including 11 features and the quality column. The quality column is an integer that ranges from 3 to 8. The goal is to build a model that predicts the quality value.

There are some problems with the raw data:

1. The column names contain space (' '), which is not compatible with Feature Store.
2. We need to add ID to the raw data so they can be looked up later by Feature Store.

The following cell addresses these issues.

Let's assume that the alcohol by volume (ABV) is a variable that changes over time after the wine is opened. The value will be provided as a real-time input for online inference.

Now, split the data into two parts and store only the part with static features to Feature Store.

Create a feature table

Save the feature data id_static_features into a feature table.

The feature data has now been stored into the feature table. The next step is to set up a Databricks Online Table.

Set up Databricks Online Tables

You can create an online table from the Catalog Explorer UI, Databricks SDK or Rest API. The steps to use Databricks python SDK are described below. For more details, see the Databricks documentation (AWS|Azure). For information about required permissions, see Permissions (AWS|Azure).

Train and deploy the model

Now, you will train a classifier using features in the Feature Store. You only need to specify the primary key, and Feature Store will fetch the required features.

First, define a TrainingSet. The training set accepts a feature_lookups list, where each item represents some features from a feature table in the Feature Store. This example uses wine_id as the lookup key to fetch all the features from table online_feature_store_example.wine_features.

The next cell trains a RandomForestClassifier model.

Save the trained model using log_model. log_model also saves lineage information between the model and the features (through training_set). So, during serving, the model automatically knows where to fetch the features by just the lookup keys.

Serve real-time queries with automatic feature lookup

After calling log_model, a new version of the model is saved.

Send a query

Now, suppose you opened a bottle of wine and you have a sensor to measure the current ABV from the bottle. Using the model and automated feature lookup with realtime serving, you can predict the quality of the wine using the measured ABV value as the realtime input "alcohol".

Notes on request format and API versions

Here is an example of the request format:

{"dataframe_split": {"index": [0, 1, 2], "columns": ["wine_id", "alcohol"], "data": [[25, 7.9], [25, 11.0], [25, 27.9]]}}

Learn more about Databricks Model Serving (AWS|Azure).

Clean up

To clean up the resources created by this notebook, follow these steps:

1. Delete the Databricks Online Table from Catalog Explorer.
   a. In the left sidebar, click Catalog.
   b. Navigate to the online table.
   c. From the kebab menu, select Delete.
2. Delete the Serving Endpoint from the Serving tab.
   a. In the left sidebar, click Serving.
   b. Click the name of the endpoint.
   c. From the kebab menu, select Delete.