Evaluating Databricks Feature Store for Machine Learning Projects
• Overview
This report evaluates Databricks Feature Store including its potential benefits to our team's ML model workflow, and its limitations. The goal is to provide a recommendation on whether our team should integrate the Feature Store capability on Databricks in our ML model development process.
• What is Databricks Feature Store?
Databricks feature store is a centralized repository designed to manage machine learning features throughout the entire lifecycle of ML models. It ensures that features are consistent and easily accessible for both training and inference, promoting reuse and collaboration among teams and projects. This integration within the Databricks platform simplifies feature management and makes model development and deployment more efficient.
• How Our Team Can Benefit From Using Databricks Feature Store
Databricks Feature Store has several key features within the Databricks ecosystem that our team can benefit from when working on the MCO project and other ML projects.
1. Easy Feature Discovery
Feature store UI in the Databricks workspace allows us to simply browse and search for existing features. This makes it easier for us to discover and reuse features, saving time and effort.
2. Feature Lineage
When creating a feature table in Databricks, it keeps track of where the data came from and how it's used. When we search for features, we can see their associated definition, source data, which models, notebooks and endpoints rely on each feature, providing transparency and traceability.
3. Integration with Model Scoring, Serving and other Databricks components
Databricks Feature Store seamlessly connects with model training and deployment. When we train a model with features from the Feature Store, it includes all relevant feature information. During batch scoring or online inference, the model automatically retrieves the required features from the Feature Store. This means we won't have to worry about adding logic to look up or link features, which simplifies the deployment process.
Databricks Feature Store streamlines feature lookups during model inference by integrating feature information into the MLflow model. This guarantees that the same transformations are used during both training and inference, lowering the possibility of inconsistencies and errors.
Databricks Feature Store seamlessly integrates with other components of the Databricks Lakehouse platform, including Delta Lake for efficient data storage and management, MLflow tracking and managing machine learning experiments, and Databricks Workflows for orchestrating complex data and ML pipelines.
4. Accurate Time-Based Lookups
Databricks Feature Store supports use cases that require features to be accurate at specific points in time, such as time series and event-based data.
5. Eliminate online/offline skew
Databricks Feature Store includes robust security features, such as access control lists (ACLs), to manage permissions. Integration with MLflow ensures that features are stored alongside ML models, facilitating governance and reducing the risk of model drift between training and serving
• Why Feature Store Is Better Than A Plain Delta Table With Transformed Features?
| Databricks Feature Store | Plain Delta Table With Transformed Features | |
|---|---|---|
| Purpose | Full ML feature lifecycle management | Simple storage of feature values |
| Feature Discoverability and Registration | Easily discoverable with UI&API, features are registered (metadata, versioning, etc) | No registration, need to manually know the table and schema to discover the features |
| Training/serving consistency | Yes | We need to manage joins and logic ourselves |
| Lineage tracking | Automatically tracks source code & transformations | Manual |
| Integration with MLflow | Automatically records feature sets in model tracking | Requires manual logging |
| Version control | Built-in versioning of feature tables | No feature-level versioning |
| Governance | Yes | Manual |
• Limitations and Considerations
- Learning curve. Team members need to spend additional time to learn the concepts, how to use feature stores, etc
- Additional effort on the coding part. Need to modify the existing code and register features.
- Potential additional compute and storage costs (limited)
• Recommendation
Based on this evaluation, adopting Databricks Feature Store is recommended if our team:
- Works on multiple ML projects that share common features / feature engineering logic
- Requires training/serving consistency and auditability
- Has long-term plans for more streamlined, scalable, production-grade ML workflows
If the features in our current MCO model is one-off and are not likely to be reused in other future ML models, then the added complexity may not be worthwhile.