What Is a Feature Store? Databricks Feature Management Guide

A Feature Store is a repository for centrally managing, sharing, and reusing the features that feed ML models. In Databricks it is delivered as "Feature Engineering in Unity Catalog", giving you a feature management platform integrated with Unity Catalog governance. Feature Store topics make up roughly 10-15% of the ML Associate exam, and questions on Feature Table creation, FeatureLookup, and Point-in-Time Join show up frequently.

Problems a Feature Store Solves

In ML projects, building features from data typically consumes the majority of development time. Without a Feature Store, the following problems frequently occur and degrade both model quality and development efficiency.

Creating a Feature Table

A Feature Table is a Delta Table with a primary key. It is managed in Unity Catalog, with the same access control and lineage tracking that apply to regular tables.

from databricks.feature_engineering import FeatureEngineeringClient

fe = FeatureEngineeringClient()

# Create a Feature Table (seed initial data from a DataFrame)
fe.create_table(
    name="catalog.schema.customer_features",
    primary_keys=["customer_id"],
    timestamp_keys=["timestamp"],       # For Point-in-Time Lookup
    df=customer_features_df,
    description="Customer behavior features (for churn prediction)"
)

# Read the features back as a DataFrame
features_df = fe.read_table(name="catalog.schema.customer_features")

primary_keys is a required parameter that uniquely identifies each row.timestamp_keys is an optional parameter used for Point-in-Time Lookup; specify it when you are working with time-series features.

Writing Features (write_table)

Use fe.write_table() to append or update features in an existing Feature Table. The mode parameter controls the write behavior.

# Add or update features in an existing Feature Table
fe.write_table(
    name="catalog.schema.customer_features",
    df=new_features_df,
    mode="merge"  # Update rows whose primary key matches; insert new rows
)

# To overwrite all data
fe.write_table(
    name="catalog.schema.customer_features",
    df=full_features_df,
    mode="overwrite"
)

• merge: rows whose primary key matches are updated (UPSERT); unmatched rows are inserted
• overwrite: replaces the entire table (used for daily batch updates)

Feature Lookup

Feature Lookup is the mechanism for assembling a training dataset from Feature Tables. Provide a DataFrame containing your labels along with the Feature Table's primary key, and the join is performed automatically.

from databricks.feature_engineering import FeatureLookup

# Define lookups from multiple Feature Tables
feature_lookups = [
    FeatureLookup(
        table_name="catalog.schema.customer_features",
        feature_names=["total_purchases", "avg_session_time", "days_since_last_visit"],
        lookup_key="customer_id"
    ),
    FeatureLookup(
        table_name="catalog.schema.product_features",
        feature_names=["category", "price_tier"],
        lookup_key="product_id"
    )
]

# Build the training dataset
training_set = fe.create_training_set(
    df=label_df,          # DataFrame containing labels and primary keys
    feature_lookups=feature_lookups,
    label="churn",
    exclude_columns=["customer_id"]  # Columns to drop from model input
)

training_df = training_set.load_df()
print(training_df.columns)
# ['total_purchases', 'avg_session_time', 'days_since_last_visit',
#  'category', 'price_tier', 'churn']

This guarantees that training and inference reference the same Feature Table, preventing Training-Serving Skew (the mismatch between training-time and serving-time features).

Point-in-Time Join

Point-in-Time Join is the mechanism that prevents data leakage when working with time-series features. It joins only the features that were actually available at the label timestamp of each training sample, keeping future information out of training.

# Configure Point-in-Time Lookup
feature_lookups = [
    FeatureLookup(
        table_name="catalog.schema.customer_features",
        feature_names=["total_purchases", "avg_session_time"],
        lookup_key="customer_id",
        timestamp_lookup_key="event_timestamp"  # Timestamp column on the label side
    )
]

# Label DataFrame (each row needs a timestamp column)
# customer_id | event_timestamp       | churn
# C001        | 2026-03-15 10:00:00   | 1
training_set = fe.create_training_set(
    df=label_df,
    feature_lookups=feature_lookups,
    label="churn"
)

In the example above, for a row whose label timestamp is 2026-03-15 10:00, only the most recent features in the Feature Table from before 2026-03-15 10:00 are joined. Features computed at 03-15 10:01 or later are not used. On the ML Associate exam, questions asking what Point-in-Time Join prevents (data leakage) are frequent.

Integrating Models with the Feature Store

When you log a model with fe.log_model(), the reference to the Feature Table (lineage) is automatically recorded in MLflow. At inference time the model automatically fetches features from the Feature Table, so callers only need to pass the primary keys.

import mlflow

# Train and log a model with the Feature Store
with mlflow.start_run():
    model = train_model(training_df)

    fe.log_model(
        model=model,
        artifact_path="model",
        flavor=mlflow.sklearn,
        training_set=training_set,
        registered_model_name="catalog.schema.churn_model"
    )

# Inference: pass a DataFrame with only primary keys; features are auto-fetched from the Feature Table
predictions = fe.score_batch(
    model_uri="models:/catalog.schema.churn_model/1",
    df=inference_df  # customer_id column is sufficient
)

The difference between fe.log_model() and mlflow.sklearn.log_model() is a frequent exam topic. Only fe.log_model() records the Feature Store linkage on the model, which is what enables automatic feature retrieval at inference time.

Online Feature Serving

Real-time inference (Model Serving) requires low-latency feature retrieval. Databricks Online Feature Serving makes Feature Table data directly available to real-time inference endpoints.

from databricks.feature_engineering.online_store_spec import (
    AmazonDynamoDBSpec
)

# Configure the online store (DynamoDB example)
online_store_spec = AmazonDynamoDBSpec(
    region="ap-northeast-1",
    table_name="customer_features_online"
)

# Publish the Feature Table to the online store
fe.publish_table(
    name="catalog.schema.customer_features",
    online_store=online_store_spec,
    mode="merge"
)

Legacy Workspace vs. Unity Catalog Version

Exam questions often ask which to use between FeatureStoreClient and FeatureEngineeringClient. For any new project, always choose the Unity Catalog version: FeatureEngineeringClient.

ML Associate Exam Focus Points

• fe.create_table(): role and configuration of primary_keys and timestamp_keys
• FeatureLookup: how to specify lookup_key and configure lookups across multiple tables
• Point-in-Time Join: definition of data leakage and how timestamp_lookup_key prevents it
• fe.log_model() vs mlflow.log_model(): difference in whether Feature Store linkage is recorded
• write_table modes: when to use merge vs. overwrite
• Training-Serving Skew: how a Feature Store solves this problem

Sample Question

Frequently Asked Questions