Feature Stores

14 minute read

“The centralized truth for machine learning features.”

1. The Problem: Feature Engineering at Scale

In ML, features are the input signals to models. Feature engineering—transforming raw data into features—often takes 80% of an ML engineer’s time.

Challenges:

Duplication: Team A and Team B independently compute the same features.
Training-Serving Skew: Features computed differently in training vs inference.
Point-in-Time Correctness: Using future data accidentally (data leakage).
Latency: Real-time features must be computed in <10ms.
Discovery: No way to know what features already exist.

Solution: A Feature Store—a centralized repository for storing, serving, and managing ML features.

2. What is a Feature Store?

A feature store is a data system that:

Stores feature values (historical and current).
Serves features for training and inference.
Manages feature metadata (schemas, lineage, freshness).
Ensures consistency between training and serving.

Key Components:

Offline Store: Stores historical features for training. Typically a data warehouse (BigQuery, Snowflake).
Online Store: Serves features for real-time inference. Typically a low-latency store (Redis, DynamoDB).
Feature Registry: Catalog of all features with metadata.
Feature Transformation Engine: Computes features from raw data.

3. Feature Store Architecture

                    ┌─────────────────────┐
                    │   Raw Data Sources  │
                    │ (Kafka, S3, DBs)    │
                    └─────────┬───────────┘
                              │
                    ┌─────────▼───────────┐
                    │ Feature Transformation│
                    │   (Spark, Flink)    │
                    └─────────┬───────────┘
                              │
         ┌────────────────────┴────────────────────┐
         │                                         │
┌────────▼────────┐                      ┌────────▼────────┐
│  Offline Store  │                      │  Online Store   │
│  (BigQuery,     │                      │  (Redis,        │
│   Snowflake)    │                      │   DynamoDB)     │
└────────┬────────┘                      └────────┬────────┘
         │                                         │
         │                                         │
┌────────▼────────┐                      ┌────────▼────────┐
│ Training Pipeline│                      │ Serving Pipeline│
│   (Batch)        │                      │   (Real-time)   │
└─────────────────┘                      └─────────────────┘

4. Offline vs Online Features

Offline Features (Batch):

Computed periodically (hourly, daily).
Used for training.
Stored in data warehouses.
Example: “Total purchases in last 30 days.”

Online Features (Real-time):

Computed at request time or updated continuously.
Used for inference.
Stored in low-latency stores.
Example: “User’s current cart value.”

Streaming Features:

Computed from event streams (Kafka, Kinesis).
Updated as events arrive.
Example: “Number of clicks in last 5 minutes.”

5. Training-Serving Skew

Problem: Features computed differently in training vs serving lead to poor model performance.

Causes:

Different Code Paths: Training uses Spark, serving uses Python.
Different Data Sources: Training uses S3, serving uses Redis.
Time Travel: Training uses historical data, serving uses current data.

Solution:

Single Definition: Define features once, compute consistently.
Feature Reuse: Serve the same precomputed features for training and inference.
Logging: Log features used in inference for debugging.

6. Point-in-Time Correctness

Problem: When training on historical data, we must use features as they existed at that time.

Example:

Target: Did user purchase on Jan 15?
Feature: User’s account age.
Correct: Account age on Jan 15.
Incorrect: Account age today (data leakage).

Solution:

Store features with timestamps.
When creating training data, join on timestamp (point-in-time join).

# Point-in-time join
def get_features_at_time(entity_id, timestamp):
    # Find the latest feature value before the timestamp
    return feature_store.query(
        entity_id=entity_id,
        timestamp_lte=timestamp  # Less than or equal
    ).order_by('timestamp', desc=True).first()

7. Feature Store Components

7.1. Feature Registry

Purpose: Catalog of all features with metadata.

Metadata:

Name: user_total_purchases_30d
Description: “Total purchases in last 30 days.”
Entity: user_id
Data Type: float
Owner: team-recommendations
Freshness: Updated daily at 2 AM.
Lineage: Computed from orders table.

Benefits:

Discovery: Find existing features before creating new ones.
Governance: Track ownership and data lineage.
Documentation: Self-documenting features.

7.2. Offline Store

Purpose: Store historical feature values for training.

Technologies:

BigQuery (GCP)
Snowflake
Redshift (AWS)
Delta Lake (Databricks)

Schema:

CREATE TABLE user_features (
    user_id STRING,
    feature_timestamp TIMESTAMP,
    total_purchases_30d FLOAT,
    avg_order_value FLOAT,
    days_since_last_login INT
);

7.3. Online Store

Purpose: Serve features with low latency for real-time inference.

Technologies:

Redis
DynamoDB (AWS)
Bigtable (GCP)
Cassandra

Requirements:

Latency: < 10ms P99.
Throughput: 100K+ reads/second.
Availability: 99.99% uptime.

Schema (Redis):

Key: user:12345:features
Value: {"total_purchases_30d": 1500.0, "avg_order_value": 75.0}
TTL: 24 hours

7.4. Feature Transformation Engine

Purpose: Compute features from raw data.

Technologies:

Apache Spark (Batch)
Apache Flink (Streaming)
dbt (SQL-based transformations)
Apache Beam

Example (Spark):

from pyspark.sql import functions as F

# Compute total purchases in last 30 days
user_features = orders_df \
    .filter(F.col("order_date") >= F.current_date() - 30) \
    .groupBy("user_id") \
    .agg(F.sum("amount").alias("total_purchases_30d"))

# Write to offline store
user_features.write.mode("overwrite").saveAsTable("user_features")

8. Popular Feature Stores

8.1. Feast (Open Source)

Features:

Lightweight, flexible.
Supports batch (offline) and online serving.
Integrates with GCP, AWS, local.

Example:

from feast import FeatureStore

store = FeatureStore(repo_path=".")

# Get training data
training_df = store.get_historical_features(
    entity_df=entity_df,  # DataFrame with entity IDs and timestamps
    features=["user_features:total_purchases_30d", "user_features:avg_order_value"]
).to_df()

# Get online features
feature_vector = store.get_online_features(
    features=["user_features:total_purchases_30d"],
    entity_rows=[{"user_id": 12345}]
).to_dict()

8.2. Tecton

Features:

Managed feature platform.
Real-time feature computation.
Strong ML Ops integration.

8.3. Databricks Feature Store

Features:

Integrated with Databricks MLflow.
Delta Lake storage.
Lineage tracking.

8.4. Amazon SageMaker Feature Store

Features:

Managed AWS service.
Online and offline stores.
Integrated with SageMaker pipelines.

8.5. Google Vertex AI Feature Store

Features:

Managed GCP service.
BigTable for online, BigQuery for offline.
AutoML integration.

9. System Design: Building a Feature Store

Scenario: Build a feature store for a ride-sharing app (like Uber).

Entities:

Driver: driver_id
Rider: rider_id
Trip: trip_id

Features:

Driver Features: rating, trips_completed, acceptance_rate, current_location.
Rider Features: rating, trips_taken, payment_method, current_location.
Real-time Features: driver’s current speed, ETA.

Architecture:

Step 1: Raw Data Sources

Event Stream (Kafka): GPS locations, trip events.
Database (PostgreSQL): User profiles, trip history.
Data Lake (S3): Historical logs.

Step 2: Feature Transformation

Batch (Spark): Compute daily aggregates (e.g., trips in last 7 days).
Streaming (Flink): Compute real-time features (e.g., speed, location).

Step 3: Storage

Offline (BigQuery): Store all historical features.
Online (Redis): Store latest features for each entity.

Step 4: Serving

Training: Query BigQuery with point-in-time joins.
Inference: Query Redis with <5ms latency.

10. Real-Time Feature Computation

Challenge: Compute features like “number of rides in last 5 minutes” in real-time.

Approach: Sliding Window Aggregations

Algorithm:

Use Apache Flink or Kafka Streams.
Define a sliding window (e.g., 5 minutes, sliding every 1 minute).
Aggregate events within the window.
Write to online store.

Example (Flink):

DataStream<TripEvent> trips = ...;

trips
    .keyBy(TripEvent::getDriverId)
    .window(SlidingEventTimeWindows.of(Time.minutes(5), Time.minutes(1)))
    .aggregate(new TripCountAggregator())
    .addSink(new RedisSink(...));

11. Feature Freshness

Definition: How often features are updated.

Trade-offs:

Fresher Features → Better Model Performance (usually).
Fresher Features → Higher Compute Cost.
Fresher Features → More Complex Infrastructure.

12. Feature Monitoring

Metrics to Track:

Freshness: How old is the latest feature value?
Coverage: What % of entities have feature values?
Distribution: Has the feature distribution shifted?
Latency: How long does it take to serve a feature?

Alerting:

Stale Features: Alert if a feature hasn’t been updated in 2x the expected period.
Distribution Drift: Alert if the feature mean/std deviates significantly.
High Latency: Alert if P99 latency > 20ms.

13. Interview Questions

What is a Feature Store? Explain its components.
Training-Serving Skew: What causes it? How do you prevent it?
Point-in-Time Correctness: Why is it important? How do you implement it?
Offline vs Online Features: When do you use each?
Design a Feature Store: For an e-commerce recommendation system.
Feature Freshness: How do you decide how often to update a feature?

14. Common Pitfalls

Over-Engineering: Starting with a complex feature store when a simple solution suffices.
Ignoring Freshness: Using stale features for real-time decisions.
No Schema Validation: Features with incorrect types or formats.
Poor Discoverability: Teams re-creating existing features.
Neglecting Monitoring: Not detecting feature drift or staleness.

15. Deep Dive: Feature Versioning

Problem: Features evolve over time. How do you manage versions?

Approach:

Semantic Versioning: user_purchases_v1, user_purchases_v2.
Schema Evolution: Support backward-compatible changes.
Deprecation Policy: Notify users before removing features.

Example:

# Feature definition with version
@feature(version=2)
def user_total_purchases(user_id):
    # V2: Include refunds
    return total_purchases - refunds

16. Deep Dive: Feature Lineage

Definition: Tracking the origin and transformations of features.

Benefits:

Debugging: Trace back to raw data when features are incorrect.
Compliance: Know which data sources are used (GDPR).
Impact Analysis: Know which models are affected when a data source changes.

Tools:

Apache Atlas
DataHub (LinkedIn)
Amundsen (Lyft)

17. Production Case Study: Uber Michelangelo

Michelangelo is Uber’s ML platform with a built-in feature store.

Scale:

Features: Millions of features.
Entities: Billions of trips, millions of users.
Latency: <10ms for online serving.

Architecture:

Offline Store: Hive.
Online Store: Cassandra.
Transformation: Spark (batch), Flink (streaming).

Key Innovations:

DSL for Feature Definitions: Single source of truth for training and serving.
Automatic Backfill: When a new feature is added, automatically compute historical values.
Feature Marketplace: Teams can share and discover features.

18. Conclusion

Feature stores are becoming essential infrastructure for ML teams. They solve the hard problems of feature management, consistency, and serving.

Key Takeaways:

Centralization: One place for all features.
Consistency: Same features for training and serving.
Point-in-Time: Avoid data leakage in training.
Freshness: Balance freshness vs cost.
Discovery: Find and reuse existing features.

The investment in a feature store pays off when:

Multiple teams use ML.
Real-time features are required.
Training-serving skew is causing issues.
Feature duplication is common.

Mastering feature stores is a key skill for ML engineers building production systems at scale.

19. Deep Dive: Embedding Features

Problem: Modern ML uses embeddings (vectors) as features. These are high-dimensional and require special handling.

Challenges:

Size: A 768-dim embedding per user = 768 * 4 bytes = 3KB per user. For 100M users = 300GB.
Indexing: Need efficient similarity search (ANN).
Freshness: Embeddings may need frequent updates.

Solutions:

Vector Databases: Pinecone, Milvus, Weaviate.
Compression: PQ (Product Quantization) reduces size by 10-100x.
Caching: Cache frequently accessed embeddings.

Example (Storing User Embeddings):

# Feast feature definition for embeddings
from feast import Entity, Feature, FeatureView, ValueType

user = Entity(name="user_id", value_type=ValueType.INT64)

user_embedding_view = FeatureView(
    name="user_embeddings",
    entities=["user_id"],
    features=[
        Feature(name="embedding", dtype=ValueType.FLOAT_LIST)
    ],
    online=True,
    batch_source=user_embedding_source
)

20. ML-Specific Data Types

Beyond Scalars:

Lists: [1.0, 2.0, 3.0] (e.g., embeddings, sequences).
Maps: {"category": 0.5, "brand": 0.3} (e.g., sparse features).
Structs: Nested objects.

Serialization:

Protocol Buffers: Efficient binary format.
Apache Arrow: Columnar format for analytics.
JSON: Human-readable but slow.

21. Cost Analysis: Build vs Buy

Build (Open Source - Feast):

Pros: Free, flexible, no vendor lock-in.
Cons: Engineering effort, maintenance, scaling.
Cost: 2-3 engineers for 6 months + ongoing maintenance.

Buy (Managed - Tecton, SageMaker):

Pros: Managed, scalable, support.
Cons: Cost, vendor lock-in.
Cost: $10K-$100K/month depending on scale.

22. Implementation Guide: From Scratch

Step 1: Define Entities

Identify primary keys: user_id, product_id, driver_id.

Step 2: Define Features

For each entity, list required features.
Classify as offline, near-real-time, or real-time.

Step 3: Choose Storage

Offline: BigQuery, Snowflake.
Online: Redis, DynamoDB.

Step 4: Build Transformation Pipelines

Batch: Spark, dbt.
Streaming: Flink, Spark Streaming.

Step 5: Create Feature Registry

Store metadata in a database (PostgreSQL).
Build a UI for discovery.

Step 6: Build Serving Layer

REST API for online serving.
Python SDK for offline retrieval.

Step 7: Add Monitoring

Track freshness, coverage, latency.
Alert on anomalies.

23. Advanced: Feature Serving Patterns

1. Precompute and Cache:

Compute features in batch, store in online store.
Serve directly from cache.
Use Case: Static features (user demographics).

2. Compute on Demand:

Compute features at request time.
Use Case: Real-time features (current session data).

3. Hybrid:

Precompute base features.
Compute derived features on demand.
Example: Precompute “total purchases”. On demand: “purchases / days since signup”.

24. Feature Engineering Best Practices

1. Avoid Data Leakage:

Never use future data as features.
Always use point-in-time joins.

2. Handle Missing Values:

Default values for missing features.
Separate “missing” indicator feature.

3. Feature Normalization:

Normalize numerical features (z-score, min-max).
Store normalization parameters in the registry.

4. Feature Documentation:

Document the business logic.
Include example values.
Link to the source data.

25. Case Study: Airbnb Zipline

Zipline is Airbnb’s feature store.

Features:

Backfill: Automatically compute historical features when a new feature is added.
Point-in-Time Joins: Built-in support for training data generation.
Streaming: Real-time features from Kafka.

Scale:

10,000+ features.
Billions of records.
Used by 100+ ML models.

26. Case Study: Spotify Feature Platform

Architecture:

Offline Store: BigQuery.
Online Store: Bigtable.
Transformation: Apache Beam.

Key Features:

Self-Service: Teams define features in YAML.
Validation: Schema validation on write.
Observability: Metrics exported to Datadog.

27. Future Trends

1. Feature Store as a Service (FSaaS):

Cloud providers offer managed feature stores.
Examples: Vertex AI Feature Store, SageMaker Feature Store.

2. Real-Time Feature Computation:

Compute features at request time using serverless functions.
Low latency, high flexibility.

3. Feature Sharing Across Organizations:

Public feature marketplaces.
Monetize features.

4. Automated Feature Engineering:

Use AutoML to generate features automatically.
Example: Featuretools, tsfresh.

28. Interview Deep Dive: Feature Store Design

Q: Design a feature store for a food delivery app (DoorDash).

Entities:

User: user_id
Restaurant: restaurant_id
Dasher (Driver): dasher_id
Order: order_id

Features:

User Features:

user_total_orders: Lifetime orders (batch, daily).
user_avg_order_value: Average order value (batch, daily).
user_current_cart: Items in cart (real-time).
user_last_location: GPS coordinates (real-time).

Restaurant Features:

restaurant_rating: Average rating (batch, daily).
restaurant_current_wait_time: Estimated wait (real-time, from Kafka).

Dasher Features:

dasher_current_location: GPS (real-time).
dasher_acceptance_rate: % of accepted orders (batch).

Architecture:

Batch Pipeline (Spark): Compute aggregates nightly.
Streaming Pipeline (Flink): Compute real-time features from GPS stream.
Offline Store (BigQuery): Historical features for training.
Online Store (Redis): Low-latency serving.

Point-in-Time Joins:

When training a model to predict delivery time, use features as they existed at order placement time.

29. Testing Feature Stores

Unit Tests:

Test feature transformation logic in isolation.
Verify expected output for sample inputs.

Integration Tests:

Test end-to-end pipeline: raw data → feature store → serving.
Verify consistency between offline and online stores.

Data Quality Tests:

Check for nulls, outliers, schema violations.
Run Great Expectations or similar tools.

30. Conclusion & Mastery Checklist

Mastery Checklist:

Feature stores are the backbone of production ML. They ensure that the features used in training are the same as those used in inference, eliminating one of the most common sources of bugs in ML systems. As ML adoption grows, feature stores will become as essential as databases are today.

Share on

Twitter Facebook LinkedIn