Python API

This page provides links to the Databricks Feature Store Python API documentation and information about the Python client package “databricks-feature-store”.

Python API reference

Use the links in the table to download or display the Feature Store Python API reference. To determine the pre-installed version for your Databricks Runtime ML version, see Feature Store compatibility matrix.

Version	Download PDF	Online API reference
v0.3.6 and above	Feature Store Python API 0.15.1 reference PDF	Online API reference
v0.3.5 and below	Feature Store Python API 0.3.5 reference PDF	Online API reference not available

Python package

The Databricks Feature Store APIs are available through the Python client package “databricks-feature-store”. The client is available on PyPI and is pre-installed in Databricks Runtime for Machine Learning. For a reference of which runtime includes which client version, see the Feature Store Compatibility Matrix.

To install the client in Databricks Runtime:

%pip install databricks-feature-store

To install the client in a local Python environment:

pip install databricks-feature-store

Supported scenarios

On Databricks, including Databricks Runtime and Databricks Runtime for Machine Learning, you can:

Create, read, and write feature tables.
Train and score models on feature data.
Publish feature tables to online stores for real-time serving.

From a local environment or an environment external to Databricks, you can:

Develop code with local IDE support.
Unit test using mock frameworks.
Write integration tests to be run on Databricks.

Limitations

The client library can only be run on Databricks, including Databricks Runtime and Databricks Runtime for Machine Learning. It does not support calling Feature Store APIs from a local environment, or from an environment other than Databricks.

Use the Feature Store client for unit testing

You can install the Feature Store client locally to aid in running unit tests.

For example, to validate that a method update_customer_features correctly calls FeatureStoreClient.write_table, you could write:

from unittest.mock import MagicMock, patch

from my_feature_update_module import update_customer_features
from databricks.feature_store import FeatureStoreClient

@patch.object(FeatureStoreClient, "write_table")
@patch("my_feature_update_module.compute_customer_features")
def test_something(compute_customer_features, mock_write_table):
  customer_features_df = MagicMock()
  compute_customer_features.return_value = customer_features_df

  update_customer_features()  # Function being tested

  mock_write_table.assert_called_once_with(
    name='recommender_system.customer_features',
    df = customer_features_df,
    mode = 'merge'
  )

Use the Feature Store client for integration testing

You can run integration tests with the Feature Store client on Databricks. For details, see Developer Tools and Guidance: Use CI/CD.

Use the Feature Store client in an integrated development environment (IDE)

You can use the Feature Store client with an IDE for software development with Databricks. For details, see Use dbx with Visual Studio Code.