How Databricks Serves Fresh ML Features in Milliseconds

Peter Bubenik · Databricks AI · · Source
Image for How Databricks Feature Store serves features with sub-second freshness

After studying this material, students should be able to:

  1. Explain why sub-second feature freshness matters in real-time ML applications
  2. Describe the end-to-end architecture of Databricks Feature Store's streaming pipeline
  3. Compare batch, microbatch, and Real-Time Mode (RTM) processing approaches
  4. Identify the role of each infrastructure component (RTM, Lakebase, Model Serving)
  5. Understand how streaming write amplification is minimized in Lakebase
  6. Recognize the three types of time windows and when to use each

Step-by-Step Study Material

Step 1: Why Feature Freshness Matters

The Core Problem

Machine learning models depend on signals (features) to make predictions. The quality of those predictions depends directly on how current those signals are.

Real-World Example: Fraud Detection

User presses "Purchase"
        ↓
Model must decide in MILLISECONDS
        ↓
Needs to know:
├── Last 30 days avg transaction amount  → Historical baseline (who is this user normally?)
└── Last 10 minutes total transactions  → Fresh signal (what are they doing RIGHT NOW?)

Key Insight: Without the fresh signal, the model cannot detect fraud happening in the moment

Why This Is Hard Today

ApproachLatencyProblem
Batch Spark jobsMinutes to hoursToo slow for real-time decisions
Custom streamingSecondsRequires complex infrastructure built by hand
Databricks Feature Store~200msManaged, declarative, automated

The Business Impact

  • Fraud detection: A transaction from 30 seconds ago could indicate an ongoing attack
  • Personalization: A user's click from 5 seconds ago reveals current intent
  • Stale features = wrong decisions = lost revenue or increased risk

Step 2: Understanding Aggregation Features and Time Windows

What Are Aggregation Features?

Aggregation features summarize raw events over a time period:

Raw events:
  Transaction $50  → 10:01am
  Transaction $200 → 10:05am
  Transaction $75  → 10:08am

Aggregation (sum, last 10 min) = $325

Common aggregation types:

  • Count — how many events occurred
  • Sum — total value of events
  • Average — mean value over the window

The Three Time Window Types

TUMBLING WINDOW (non-overlapping, fixed)
|--10min--|--10min--|--10min--|
  Batch 1   Batch 2   Batch 3

SLIDING WINDOW (overlapping, fixed step)
|--10min--|
    |--10min--|
        |--10min--|

ROLLING WINDOW (continuous, event-driven)
Every new event immediately updates the aggregate
New event → instant recalculation → instant update

When to Use Each

Window TypeBest ForUpdate FrequencyCost
TumblingStable, infrequent featuresLowCheapest
SlidingModerate freshness needsMediumModerate
RollingMaximum freshnessEvery eventHigher

Rule of thumb: Use rolling windows when every new event should immediately change the value the model sees

Defining a Rolling Window Feature (Declarative API)

The Feature Store lets you define this simply:

# Conceptual example of the declarative API
feature = RollingWindowFeature(
    name="user_transaction_sum_10min",
    entity="user_id",
    aggregation="sum",
    column="transaction_amount",
    window="10 minutes"
)

One definition → drives both batch and streaming pipelines automatically


Step 3: The End-to-End Streaming Pipeline Architecture

Overview of the Pipeline

Kafka (raw events)
      ↓
  Spark RTM (Real-Time Mode)
  ├── Stage 1: Data Processing
  │     └── Schema validation, type casting, coalescing
  ├── Stage 2: Aggregation
  │     └── RocksDB (local state) → rolling window calculation
      ↓
  Lakebase (online feature store)
      ↓
  Model Serving (inference)
      ↓
  ML Model Decision

Fraud Detection Example Traced Through the Pipeline

Event arrives: User buys $200 item
        ↓
Kafka receives event with:
  - user_id: 12345
  - amount: $200
  - timestamp: now
  - location, merchant info...
        ↓
RTM Stage 1: Validate and shape the data
        ↓
RTM Stage 2: 
  - Look up user 12345 in local RocksDB
  - Current running total: $125 (from last 10 min)
  - Add $200 → new total: $325
  - Remove expired events (older than 10 min)
  - Write $325 to Lakebase
        ↓
Model query arrives: "Should I approve this transaction?"
  - Fetch fresh feature: $325 (10-min sum)
  - Fetch batch feature: $45 avg (30-day baseline)
  - $325 >> $45 baseline → HIGH FRAUD RISK
        ↓
Decision: DECLINE

Step 4: Spark Real-Time Mode (RTM) — The Key Innovation

The Problem with Traditional Microbatch Mode (MBM)

MICROBATCH MODE:
─────────────────────────────────────────────────
Collect events → Process batch → Checkpoint → Collect events → ...
   (wait)           (wait)         (wait)        (wait)

Minimum latency: SECONDS to MINUTES
Each stage must FINISH before the next begins
─────────────────────────────────────────────────

How RTM Is Different

REAL-TIME MODE:
─────────────────────────────────────────────────
Stage 1 ──→ processes row immediately
              ↓ (passes row forward without waiting)
Stage 2 ──→ aggregates row immediately
              ↓
Lakebase ──→ updated immediately

Stages run CONCURRENTLY, not sequentially
─────────────────────────────────────────────────

Analogy: MBM is like a factory that waits until 100 items are built before moving them to the next station. RTM moves each item to the next station the moment it's ready.

RocksDB: Local State Storage

Each executor maintains a local RocksDB instance:

Why RocksDB?
├── Runs locally on each executor (no network round-trip)
├── Can store state LARGER than available memory (disk-backed)
├── Fast key-value lookups by entity (e.g., user_id)
└── Handles window expiration per-row

Checkpointing: Fault Tolerance Without Sacrificing Speed

AspectMicrobatchRTM
When checkpointedEvery batch boundaryEvery 5 minutes
Effect on latencyBlocks pipeline each timeCost spread across many rows
Fault toleranceFullFull (replays ≤5 min of Kafka data)
Recovery costLower replaySlightly higher replay

Key tradeoff: RTM accepts slightly more replay data on failure in exchange for dramatically lower steady-state latency

Serverless Deployment via Lakeflow SDP

You don't manage:
  ✗ Cluster provisioning
  ✗ Executor tuning
  ✗ Capacity planning
  ✗ Maintenance windows

Databricks handles:
  ✓ Auto-scaling
  ✓ Zero-downtime restarts (synchronized at 5-min checkpoints)
  ✓ Infrastructure updates

Step 5: Lakebase — Streaming-Optimized Online Storage

What Is Lakebase?

Lakebase is Databricks' online feature store — a Postgres-compatible database optimized for:

  • High-throughput streaming writes
  • Low-latency reads (tens of milliseconds)
  • Tens of thousands of reads per second

The Write Amplification Problem (Standard Postgres)

Standard Postgres WAL behavior:
─────────────────────────────────────────────────
Event updates user 12345's feature value
        ↓
Postgres writes FULL 8KB page to WAL
(even though only a few bytes changed)
        ↓
Next update to same user → ANOTHER full 8KB page
        ↓
High-frequency updates = MASSIVE WAL bloat
= Bottleneck for writes, replication, recovery
─────────────────────────────────────────────────

How Lakebase Solves This

Lakebase architecture:
─────────────────────────────────────────────────
Event updates user 12345's feature value
        ↓
Lakebase writes SMALL compact change record
(only the actual change, not full page)
        ↓
Quorum of distributed safekeeper nodes
acknowledge durability
        ↓
Full page snapshots generated LATER in storage layer
(not in the write path)
─────────────────────────────────────────────────
Result: Far less WAL amplification
        Continuous fresh writes possible
        Durability still guaranteed

Compute-Storage Separation Benefits

Lakebase
├── Compute layer: Scales independently
│     └── Handles variable inference load
└── Storage layer: Distributed, durable
      └── Handles streaming write volume

Result: Autoscales to handle spikes in model inference
        without affecting write performance

Step 6: Model Serving — Delivering Features at Inference Time

How Model Serving Works with Feature Store

Model logged with MLflow
        ↓
Feature dependencies automatically recorded
        ↓
At inference time:
  Request arrives → Model Serving automatically:
    1. Looks up required features from Lakebase
    2. Joins features with inference request
    3. Passes complete feature vector to model
    4. Returns prediction
        ↓
No custom lookup code needed
No manual plumbing

Performance Characteristics

  • Optimized for high QPS (queries per second)
  • Optimized for low latency
  • Fully managed — no infrastructure to operate

Step 7: Two Additional Capabilities

Training Data Generation

The problem: Streaming features use short retention windows on Kafka — historical data disappears

The solution:

Databricks Feature Store:
├── Stores offline copy of ingested Kafka data
├── Calculates same feature values as streaming pipeline
│     (for historical periods)
└── Does point-in-time accurate joins for training
      (no data leakage — only uses data available at that moment)

Also enables: Fast backfilling of online features for new deployments

Governance and Lineage via Unity Catalog

Features are first-class objects:
├── Discoverable (searchable catalog)
├── Access-controlled (who can use which features)
├── Lineage tracked (which model uses which features)
└── MLflow captures feature-to-model dependencies

Result: One platform for develop → deploy → govern

Summary: The Complete Picture

┌─────────────────────────────────────────────────────────┐
│              Databricks Feature Store                    │
│                                                         │
│  Define once → Use everywhere                           │
│                                                         │
│  Kafka → [Spark RTM] → [Lakebase] → [Model Serving]    │
│                                                         │
│  End-to-end p99 latency: 200ms                         │
└─────────────────────────────────────────────────────────┘
ComponentRoleKey Optimization
Spark RTMContinuous stream processingConcurrent stages, per-row aggregation
RocksDBLocal state for aggregationsNo network latency, disk-backed
LakebaseOnline feature storageCompact WAL writes, compute-storage separation
Model ServingFeature retrieval at inferenceAutomatic feature lookup, high QPS
Unity CatalogGovernance and lineageDiscoverable, access-controlled features

Knowledge Check Questions

  1. Why is a 30-day batch feature insufficient alone for fraud detection?
  2. What is the fundamental difference between microbatch and Real-Time Mode processing?
  3. Why does standard Postgres struggle with high-frequency streaming writes?
  4. What tradeoff does RTM make regarding checkpointing and fault tolerance?
  5. How does Databricks Feature Store solve the training data problem for streaming features?
  6. When would you choose a tumbling window over a rolling window?

More to study