Online Feature Store example with Unity Catalog

This notebook illustrates the use of Databricks Feature Engineering in Unity Catalog to publish features to an online store 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 Cosmos DB
   • This notebook uses Cosmos DB. For a list of supported online stores, see the Databricks documentation.
3. Publish the features to online feature store
4. Train and deploy the model
5. Serve realtime queries with automatic feature lookup
6. Clean up

Data Set

This example uses the Wine Quality Data Set.

Requirements

• Databricks Runtime 13.3 LTS for Machine Learning or above.
  • If you do not have access to Databricks Runtime for Machine Learning, you can run this notebook on Databricks Runtime 13.3 LTS or above. To do so, run %pip install databricks-feature-engineering at the start of this notebook.
• Access to Azure Cosmos DB.
  • This notebook uses Cosmos DB as the online store and guides you through how to generate secrets and register them with Databricks Secret Management.
• The cluster you are running must have the Azure Cosmos DB connector for Spark installed. See the instructions in the section Prepare the compute cluster.

Prepare the compute cluster

1. When creating the compute cluster, select Unrestricted or Shared Compute policy. To run this notebook on a Shared Compute cluster, you must select Databricks Runtime for ML 13.2 or above.
2. After creating the cluster, follow these steps to install the latest Azure Cosmos DB connector for Spark 3.3:
   • Navigate to the compute cluster and click the Libraries tab.
   • Click Install new.
   • Click Maven and enter the coordinates of the latest version. For example: com.azure.cosmos.spark:azure-cosmos-spark_3-3_2-12:4.20.0
   • Click Install.
3. Attach this notebook to the cluster.

Prepare the feature table

Suppose you need to build an endpoint to predict wine quality with just the wine_id. There has to be a feature table saved in Feature Engineering in Unity Catalog 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 Engineering in Unity Catalog.
2. We need to add ID to the raw data so they can be looked up later by Feature Engineering in Unity Catalog.

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 realtime input in online inference.

Now, split the data into two parts and store only the part with static features to Feature Engineering in Unity Catalog.

Create a feature table

Next, create a new catalog or reuse an existing one and create the schema where the feature tables will be stored.

• To create a new catalog, you must have the CREATE CATALOG privilege on the metastore.
• To use an existing catalog, you must have the USE CATALOG privilege on the catalog.
• To create a new schema in the catalog, you must have the CREATE SCHEMA privilege on the catalog.

Save the feature data id_static_features into a feature table.

The feature data has been stored into the feature table. The next step is to set up access to Azure Cosmos DB.

Set up Cosmos DB credentials

In this section, you need to take some manual steps to make Cosmos DB accessible to this notebook. Databricks needs permission to create and update Cosmos DB containers so that Cosmos DB can work with Feature Engineering in Unity Catalog. The following steps stores Cosmos DB keys in Databricks Secrets.

Look up the keys for Cosmos DB

1. Go to Azure portal at https://portal.azure.com/
2. Search and open "Cosmos DB", then create or select an account.
3. Navigate to "keys" the view the URI and credentials.

Provide online store credentials using Databricks secrets

Note: For simplicity, the commands below use predefined names for the scope and secrets. To choose your own scope and secret names, follow the process in the Databricks documentation.

1. Create two secret scopes in Databricks.

   databricks secrets create-scope --scope feature-store-example-read
   databricks secrets create-scope --scope feature-store-example-write
   

2. Create secrets in the scopes. Note: the keys should follow the format <prefix>-authorization-key. For simplicity, these commands use predefined names here. When the commands run, you will be prompted to copy your secrets into an editor.

   databricks secrets put --scope feature-store-example-read --key cosmos-authorization-key
   databricks secrets put --scope feature-store-example-write --key cosmos-authorization-key
   

Now the credentials are stored with Databricks Secrets. You will use them below to create the online store.

Publish the features to the online feature store

This allows Feature Engineering in Unity Catalog to add a lineage information about the feature table and the online storage. So when the model serves real-time queries, it can lookup features from the online store for better performance.

Note: You must use publish_table() to create the table in the online store. Do not manually create a database or container inside Cosmos DB. publish_table() does that for you automatically. If you create a table without using publish_table(), the schema might be incompatible and the write command will fail.

Train and deploy the model

Now, you will train a classifier using features in Feature Engineering in Unity Catalog. You only need to specify the primary key, and Feature Engineering in Unity Catalog 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 Feature Engineering in Unity Catalog. This example uses wine_id as the lookup key to fetch all the features from table ml.online_store_example.wine_static_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 realtime queries with automatic feature lookup

After calling log_model, a new version of the model is saved. To provision a serving endpoint, follow the steps below.

1. Click Serving under Machine Learning in the left sidebar.
2. Create a serving endpoint with the model named "wine_quality_classifier". See the Databricks documentation for details.

Send a query

In the Serving page, there are three approaches for calling the model. You can try the "Browser" approach with a JSON format request, as shown below. But here we copy-pasted the Python approach to illustrate an programatic way.

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.

Clean up

Follow this checklist to clean up the resources created by this notebook:

1. Azure Cosmos DB Container
   • Go to Azure console and navigate to Cosmos DB.
   • Delete the container feature_store_online_wine_features
2. Secrets store on Databricks Secrets databricks secrets delete-scope --scope <scope-name>
3. Databricks access token
   • Click your username at the top right. Then select User Settings > Access tokens and delete the token.