How do I invoke a data enrichment function before model.predict while serving the model in Databricks

Question

In Databricks, I have used mlflow and got my model served through REST API. It works fine when all model features are provided. But my use case is that only a single feature (the primary key) will be provided by the consumer application, and my code has to lookup the other features from a database based on that key and then use the model.predict to return the prediction. I tried researching but understood that the REST endpoints will simply invoke the model.predict function. How can I make it invoke a data massaging function before predicting?

Answer 1

There are two approaches for that:

1. You can use custom MLflow model, where you override the predict function, and it will call database or other source for an additional information, and then call actual predict of the model. You can find more information in following answers: 1, 2.

2. Use Databricks Feature Store for your data, train & log model using the FeatureStoreClient.log_model function, then publish feature store tables into a database, and then use model via model serving, and it will automatically lookup for features.