DEV Community: Jubin Soni

Spark Performance Deep Dive on Databricks: Shuffle Tuning, Skew Handling, and Z-Ordering with Delta Lake + Unity Catalog

Jubin Soni — Mon, 27 Jul 2026 03:14:33 +0000

The problem with "just add more workers"

Most Spark performance issues on Databricks aren't solved by scaling the cluster — they're caused by shuffle and skew, and no amount of extra nodes fixes a badly partitioned join. This post builds a realistic pipeline (order events joined against a small dimension table, aggregated, and written to Delta Lake) from the ground up, and uses it to work through:

How Spark's shuffle actually behaves during a wide transformation
Diagnosing and fixing data skew with salting and adaptive query execution (AQE)
Laying out the resulting Delta table with Z-Ordering so downstream queries skip irrelevant files
Governing access to the whole pipeline with Unity Catalog

Architecture overview

Pipeline shape — a batch job reading raw events, joining against a dimension table, aggregating, and writing to a governed Delta table:

What happens inside a shuffle stage — this is the part most tutorials skip, and it's the key to understanding why skew hurts:

Step 1 — Set up governed tables in Unity Catalog

Everything downstream depends on tables being registered under Unity Catalog, which gives you centralized access control and lineage instead of per-workspace table grants.

-- setup.sql, run in a Databricks SQL or notebook cell
CREATE CATALOG IF NOT EXISTS retail_analytics;
CREATE SCHEMA IF NOT EXISTS retail_analytics.events;

CREATE TABLE IF NOT EXISTS retail_analytics.events.raw_orders (
  order_id STRING,
  customer_id STRING,
  product_id STRING,
  quantity INT,
  event_ts TIMESTAMP
)
USING DELTA
LOCATION 'abfss://data@<storage-account>.dfs.core.windows.net/raw_orders';

CREATE TABLE IF NOT EXISTS retail_analytics.events.dim_products (
  product_id STRING,
  category STRING,
  unit_cost DOUBLE
)
USING DELTA
LOCATION 'abfss://data@<storage-account>.dfs.core.windows.net/dim_products';

GRANT SELECT ON TABLE retail_analytics.events.raw_orders TO `analysts`;

Step 2 — Read and force a broadcast join for the small dimension table

dim_products is small, so let Spark broadcast it rather than shuffle both sides of the join.

# pipeline.py
from pyspark.sql import functions as F

orders = spark.table("retail_analytics.events.raw_orders")
products = spark.table("retail_analytics.events.dim_products")

joined = orders.join(F.broadcast(products), on="product_id", how="left")

Without the explicit broadcast hint, Spark's cost-based optimizer usually picks a broadcast join automatically for small tables, but being explicit avoids surprises when the dimension table grows past spark.sql.autoBroadcastJoinThreshold (default 10MB) without anyone noticing.

Step 3 — The aggregation that triggers a shuffle

groupBy on customer_id is a wide transformation — Spark must shuffle rows so all records for a given key land on the same reducer.

agg = (
    joined
    .groupBy("customer_id", "category")
    .agg(
        F.sum(F.col("quantity") * F.col("unit_cost")).alias("total_spend"),
        F.count("order_id").alias("order_count"),
    )
)

If one customer_id (say, a test account or a large B2B buyer) accounts for a disproportionate share of rows, this is where skew shows up: one reducer task runs for minutes while the rest of the stage finishes in seconds. You'll see this in the Spark UI as a single long-running task in an otherwise short stage.

Step 4 — Fixing skew with salting

Adaptive Query Execution (AQE) handles a lot of skew automatically in modern Databricks Runtime, but for known hot keys, explicit salting is still the most predictable fix.

from pyspark.sql import functions as F
import random

SALT_BUCKETS = 20

# Add a salt column to spread the hot key across multiple reducers
salted = joined.withColumn("salt", (F.rand() * SALT_BUCKETS).cast("int"))

partial_agg = (
    salted
    .groupBy("customer_id", "category", "salt")
    .agg(
        F.sum(F.col("quantity") * F.col("unit_cost")).alias("total_spend"),
        F.count("order_id").alias("order_count"),
    )
)

# Second pass: combine the salted partial aggregates into the final result
final_agg = (
    partial_agg
    .groupBy("customer_id", "category")
    .agg(
        F.sum("total_spend").alias("total_spend"),
        F.sum("order_count").alias("order_count"),
    )
)

This two-phase pattern — pre-aggregate on a salted key, then combine — is the same trick used inside combiners in older MapReduce systems. It trades a bit of extra shuffle for eliminating the single-reducer bottleneck.

Also worth setting explicitly rather than relying on the default of 200:

spark.conf.set("spark.sql.shuffle.partitions", "auto")  # let AQE size it dynamically
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

Step 5 — Write to Delta Lake with optimized writes

(
    final_agg.write
    .format("delta")
    .mode("overwrite")
    .option("delta.autoOptimize.optimizeWrite", "true")
    .saveAsTable("retail_analytics.events.customer_spend_summary")
)

Optimized writes shuffle data before writing so you get fewer, better-sized files instead of many small ones — this costs some write-time latency in exchange for faster reads later.

Step 6 — Z-Order the table for the queries that matter

If most downstream queries filter on customer_id, co-locate related rows physically so Spark can skip files that can't match the filter.

OPTIMIZE retail_analytics.events.customer_spend_summary
ZORDER BY (customer_id);

Z-ordering aims to produce evenly balanced data files by row count rather than raw size, and its effectiveness depends on the column having reasonably high cardinality — Z-ordering by a low-cardinality column like category alone gives little benefit. On newer Delta Lake / Databricks Runtime versions, Liquid Clustering is generally the preferred choice for new tables since it adapts as query patterns change, while ZORDER remains relevant mainly for existing tables not yet migrated.

-- Preferred on new tables (Databricks Runtime supporting Liquid Clustering):
CREATE TABLE retail_analytics.events.customer_spend_summary
CLUSTER BY (customer_id);

Comparing shuffle-mitigation techniques

Technique	Fixes	Cost	When to use
Broadcast join	Shuffle on the large side of a join	Extra memory on executors	Small (<~10MB by default) dimension/lookup tables
AQE skew join handling	Automatic detection of skewed partitions	Minor planning overhead	Default-on; good general safety net
Manual salting	Known, severe hot keys	Extra shuffle for the two-phase aggregate	Hot keys that AQE doesn't fully resolve
Repartition by key	Uneven task distribution before a shuffle stage	One extra shuffle	Pre-shaping data before multiple downstream joins
Z-Ordering	Slow reads due to unnecessary file scans	Shuffle + rewrite during OPTIMIZE	Existing tables, high-cardinality filter columns
Liquid Clustering	Same as Z-Order, plus evolving query patterns	Ongoing background clustering	New tables on supporting runtime versions

Governance: tying it back to Unity Catalog

Because every table above was created under retail_analytics, access control, audit logging, and lineage are handled centrally rather than per-cluster:

-- Restrict PII-adjacent columns without duplicating the table
CREATE VIEW retail_analytics.events.customer_spend_summary_masked AS
SELECT
  customer_id,
  category,
  total_spend,
  order_count
FROM retail_analytics.events.customer_spend_summary;

GRANT SELECT ON VIEW retail_analytics.events.customer_spend_summary_masked TO `bi_readers`;

Production considerations

Diagnose before tuning. Check the Spark UI's stage view for one long-running task among many short ones — that's the signature of skew, not just "the job is slow."
Predictive optimization can run OPTIMIZE and ANALYZE automatically on Unity Catalog managed tables, which reduces the need for scheduled maintenance jobs for many workloads.
Don't Z-Order everything. It requires shuffling and rewriting the table (or partition), so reserve it for columns that are actually filtered on frequently downstream.

References

Grounded Search Agents: Microsoft Foundry + Azure AI Search + Cosmos DB + Function Calling

Jubin Soni — Mon, 27 Jul 2026 01:58:15 +0000

Why "search" isn't enough anymore

A plain RAG pipeline — embed a query, hit a vector index, stuff the top-k chunks into a prompt — works fine for a single-turn FAQ bot. It falls apart the moment you need:

Multi-turn memory ("what about last quarter?" needs to know what "last quarter" refers to)
A decision about whether to search at all (some questions are structured lookups, not document retrieval)
Citations the caller can verify

This post builds a small but complete grounded search agent: Azure OpenAI (via Microsoft Foundry) decides when to call a search tool versus a structured data tool, Azure AI Search handles retrieval and semantic ranking over unstructured documents, and Azure Cosmos DB stores conversation state and structured order/customer metadata. The agent uses OpenAI function calling to route between the two data sources and always returns citations.

We'll build this from the ground up in Python.

Architecture overview

The system has two phases: ingestion (getting documents into a searchable index) and runtime (the agent loop that answers questions).

At runtime, the agent sits between the user and two tools, using function calling to decide which one (or both) to invoke:

Step 1 — Provision and index documents in Azure AI Search

Create the index with both a vector field (for semantic similarity) and a semantic configuration (for ranking), since combining the two consistently outperforms either alone.

# create_index.py
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SimpleField, SearchableField, SearchField,
    SearchFieldDataType, VectorSearch, HnswAlgorithmConfiguration,
    VectorSearchProfile, SemanticConfiguration, SemanticPrioritizedFields,
    SemanticField, SemanticSearch
)
from azure.core.credentials import AzureKeyCredential

endpoint = "https://<your-search-service>.search.windows.net"
credential = AzureKeyCredential("<admin-key>")
index_client = SearchIndexClient(endpoint, credential)

fields = [
    SimpleField(name="id", type=SearchFieldDataType.String, key=True),
    SearchableField(name="content", type=SearchFieldDataType.String),
    SimpleField(name="source_url", type=SearchFieldDataType.String),
    SearchField(
        name="content_vector",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=1536,
        vector_search_profile_name="default-profile",
    ),
]

vector_search = VectorSearch(
    profiles=[VectorSearchProfile(name="default-profile", algorithm_configuration_name="hnsw-config")],
    algorithms=[HnswAlgorithmConfiguration(name="hnsw-config")],
)

semantic_config = SemanticConfiguration(
    name="default-semantic",
    prioritized_fields=SemanticPrioritizedFields(
        content_fields=[SemanticField(field_name="content")]
    ),
)

index = SearchIndex(
    name="support-docs-index",
    fields=fields,
    vector_search=vector_search,
    semantic_search=SemanticSearch(configurations=[semantic_config]),
)

index_client.create_or_update_index(index)
print("Index created.")

Now embed and upload chunks. Chunking matters more than most tutorials admit — 512-token chunks with ~15% overlap is a reasonable default for policy/support documents.

# ingest.py
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
import uuid

aoai = AzureOpenAI(
    azure_endpoint="https://<your-foundry-resource>.openai.azure.com",
    api_key="<key>",
    api_version="2024-10-21",
)
search_client = SearchClient(
    endpoint=endpoint,
    index_name="support-docs-index",
    credential=AzureKeyCredential("<admin-key>"),
)

def chunk_text(text: str, max_tokens: int = 512, overlap: int = 80):
    words = text.split()
    step = max_tokens - overlap
    return [" ".join(words[i:i + max_tokens]) for i in range(0, len(words), step)]

def embed(text: str) -> list[float]:
    resp = aoai.embeddings.create(model="text-embedding-3-small", input=text)
    return resp.data[0].embedding

def ingest_document(text: str, source_url: str):
    docs = []
    for chunk in chunk_text(text):
        docs.append({
            "id": str(uuid.uuid4()),
            "content": chunk,
            "source_url": source_url,
            "content_vector": embed(chunk),
        })
    search_client.upload_documents(documents=docs)

with open("return_policy.txt") as f:
    ingest_document(f.read(), source_url="policies/return_policy.txt")

Step 2 — Structured lookups in Cosmos DB

Not every question is a document-retrieval problem. "Where's my order?" is a point lookup. Store operational data in Cosmos DB, partitioned by customer_id for predictable query performance.

# cosmos_setup.py
from azure.cosmos import CosmosClient, PartitionKey

client = CosmosClient("<cosmos-endpoint>", credential="<cosmos-key>")
db = client.create_database_if_not_exists("retail")
orders = db.create_container_if_not_exists(
    id="orders",
    partition_key=PartitionKey(path="/customer_id"),
    offer_throughput=400,
)

orders.upsert_item({
    "id": "4471",
    "customer_id": "cust_9001",
    "status": "delayed",
    "delay_reason": "carrier weather disruption",
    "eta": "2026-08-02",
})

Step 3 — The agent loop with function calling

This is the core piece: Azure OpenAI decides, per turn, whether to call search_docs, get_order, both, or neither. We define the tools as JSON schemas and let the model pick.

# agent.py
import json
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.cosmos import CosmosClient

aoai = AzureOpenAI(
    azure_endpoint="https://<your-foundry-resource>.openai.azure.com",
    api_key="<key>",
    api_version="2024-10-21",
)
search_client = SearchClient(endpoint, "support-docs-index", AzureKeyCredential("<admin-key>"))
cosmos = CosmosClient("<cosmos-endpoint>", credential="<cosmos-key>")
orders_container = cosmos.get_database_client("retail").get_container_client("orders")

tools = [
    {
        "type": "function",
        "function": {
            "name": "search_docs",
            "description": "Semantic search over policy and support documents. Returns passages with citations.",
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string"}},
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_order",
            "description": "Look up a customer order by ID for status, ETA, and delay reason.",
            "parameters": {
                "type": "object",
                "properties": {
                    "order_id": {"type": "string"},
                    "customer_id": {"type": "string"},
                },
                "required": ["order_id", "customer_id"],
            },
        },
    },
]

def search_docs(query: str):
    results = search_client.search(
        search_text=query,
        query_type="semantic",
        semantic_configuration_name="default-semantic",
        top=3,
    )
    return [{"content": r["content"], "source": r["source_url"]} for r in results]

def get_order(order_id: str, customer_id: str):
    item = orders_container.read_item(item=order_id, partition_key=customer_id)
    return {"status": item["status"], "reason": item.get("delay_reason"), "eta": item.get("eta")}

def run_agent(user_message: str, customer_id: str, history: list):
    messages = history + [{"role": "user", "content": user_message}]
    response = aoai.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tools,
        tool_choice="auto",
    )
    msg = response.choices[0].message

    if msg.tool_calls:
        messages.append(msg)
        for call in msg.tool_calls:
            args = json.loads(call.function.arguments)
            if call.function.name == "search_docs":
                result = search_docs(**args)
            elif call.function.name == "get_order":
                args["customer_id"] = customer_id  # enforce, don't trust model-provided id
                result = get_order(**args)
            messages.append({
                "role": "tool",
                "tool_call_id": call.id,
                "content": json.dumps(result),
            })
        final = aoai.chat.completions.create(model="gpt-4o", messages=messages)
        return final.choices[0].message.content

    return msg.content

A few things worth calling out:

Never trust model-supplied identity fields. customer_id is overwritten server-side before hitting Cosmos DB, even though the model could technically supply one — this prevents cross-tenant data leaks via prompt injection.
tool_choice="auto" lets the model skip tools entirely for small talk, avoiding unnecessary search calls.
The second chat.completions.create call (after tool results are appended) is what produces the grounded, cited final answer.

Step 4 — Session state in Cosmos DB

Conversation history should not live in application memory — store it in Cosmos DB so any instance of the service can resume a session.

# session_store.py
def load_history(session_id: str) -> list:
    container = cosmos.get_database_client("retail").get_container_client("sessions")
    try:
        item = container.read_item(item=session_id, partition_key=session_id)
        return item["messages"]
    except Exception:
        return []

def save_history(session_id: str, messages: list):
    container = cosmos.get_database_client("retail").get_container_client("sessions")
    container.upsert_item({"id": session_id, "messages": messages})

Comparing retrieval strategies

Choosing between search modes is a real architectural decision, not just a config toggle:

Retrieval Mode	Best for	Weakness
Keyword (BM25)	Exact terms, SKUs, error codes	Misses paraphrased queries
Pure vector search	Conceptual / paraphrased questions	Can surface plausible-but-wrong matches
Hybrid (keyword + vector)	General-purpose document Q&A	Slightly higher latency than either alone
Semantic ranking on top of hybrid	Conversational, nuanced queries	Adds a reranking pass; not needed for simple lookups
Structured query (Cosmos DB point read)	Known-ID lookups (orders, accounts)	Not applicable to unstructured content

For most support/document agents, hybrid retrieval with semantic ranking is the right default, reserving structured lookups for anything with a stable ID.

Production considerations

Grounding failures are silent failures. If retrieval returns weak matches, the model can still generate a fluent but wrong answer — always surface citations so users (and evaluators) can catch this.
Access control matters. If your index mixes content from different permission levels, filter search results by the caller's identity before they reach the model, not after.
Rising input tokens. Retrieved passages add to the prompt on every call — budget for this in both cost and latency.

References

Building Production-Grade Semantic Search with GPT-5 and Microsoft Foundry, From Scratch

Jubin Soni — Sun, 19 Jul 2026 03:40:11 +0000

Most "RAG tutorials" stop at a single embedding query against a single index. That works for a demo and falls over the moment a real user asks something like "compare our Q3 and Q4 vendor contracts and flag anything that changed" — a question that needs multiple sub-queries, reasoning about what's still missing, and synthesis across documents.

Microsoft Foundry's answer to this is Foundry IQ: an agentic retrieval layer built on Azure AI Search that treats retrieval as a reasoning task rather than a single keyword or vector lookup. This is a from-scratch build of a semantic search pipeline using GPT-5 for query planning/synthesis and Foundry IQ for retrieval.

Architecture

Instead of one query hitting one index once, Foundry IQ's knowledge base plans sub-queries, executes them in parallel against one or more knowledge sources, evaluates whether it has enough signal, and iterates before synthesizing a final, cited answer.

The knowledge base sits between your agent and the underlying content. Your Foundry agent doesn't talk to Azure AI Search directly — it calls the knowledge base's MCP endpoint, which handles planning, retrieval, and synthesis behind a single tool call.

Step 1: Create a knowledge source

A knowledge source is a reusable reference to your underlying content — in this example, a Blob Storage container of documents. Creating it also triggers Azure AI Search to generate the index, skillset, and indexer needed to chunk and vectorize the content automatically.

from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndexKnowledgeSource,
    SearchIndexKnowledgeSourceParameters,
)
from azure.identity import DefaultAzureCredential

index_client = SearchIndexClient(
    endpoint="https://<your-search-service>.search.windows.net",
    credential=DefaultAzureCredential(),
)

knowledge_source = SearchIndexKnowledgeSource(
    name="vendor-contracts-ks",
    search_index_parameters=SearchIndexKnowledgeSourceParameters(
        search_index_name="vendor-contracts-index",
    ),
)
index_client.create_or_update_knowledge_source(knowledge_source=knowledge_source)

Step 2: Create a knowledge base with GPT-5 for planning and synthesis

The knowledge base ties one or more knowledge sources together with an LLM deployment that handles query planning and answer synthesis.

from azure.search.documents.indexes.models import (
    KnowledgeBase,
    KnowledgeBaseAzureOpenAIModel,
    AzureOpenAIVectorizerParameters,
)

knowledge_base = KnowledgeBase(
    name="vendor-contracts-kb",
    knowledge_sources=[{"name": "vendor-contracts-ks"}],
    models=[
        KnowledgeBaseAzureOpenAIModel(
            azure_open_ai_parameters=AzureOpenAIVectorizerParameters(
                resource_url="https://<your-foundry-resource>.openai.azure.com",
                deployment_name="gpt-5-mini",
                model_name="gpt-5-mini",
            )
        )
    ],
    output_configuration={
        "modality": "answerSynthesis",  # verbatim extractive data is the alternative
    },
)
index_client.create_or_update_knowledge_base(knowledge_base=knowledge_base)

output_configuration is the key lever here: answerSynthesis returns a pre-generated, cited answer, while extractive mode returns verbatim source chunks and leaves reasoning entirely to your agent's own model. Extractive mode costs less and gives your agent more control; synthesis mode does more work up front at the retrieval layer.

Step 3: Wire the knowledge base into a Foundry agent via MCP

Each knowledge base exposes a standalone MCP endpoint. Any MCP-compatible client — including Foundry Agent Service, but also GitHub Copilot or other MCP clients — can call its knowledge_base_retrieve tool directly.

from azure.ai.projects import AIProjectClient
from azure.identity import DefaultAzureCredential

project = AIProjectClient(
    endpoint="https://<your-foundry-project-endpoint>",
    credential=DefaultAzureCredential(),
)

agent = project.agents.create_agent(
    model="gpt-5-mini",
    name="contract-search-agent",
    instructions="Answer questions about vendor contracts using the knowledge base tool. Always cite sources.",
    tools=[{
        "type": "mcp",
        "server_url": "https://<your-search-service>.search.windows.net/knowledgebases/vendor-contracts-kb/mcp?api-version=2026-05-01-preview",
        "allowed_tools": ["knowledge_base_retrieve"],
    }],
)

Reasoning effort and cost/latency tradeoffs

The knowledge base's retrieval_reasoning_effort setting controls how much LLM-driven planning happens before retrieval, and it's the main dial for balancing latency, cost, and answer quality.

Reasoning effort	What happens	Best for
Minimal	Bypasses LLM query planning entirely; direct hybrid search	Simple factual lookups, latency-sensitive paths
Medium	LLM reformulates and may decompose the query into sub-queries	Multi-part or ambiguous questions
High	Full iterative planning, evaluates sufficiency, re-queries as needed	Complex, multi-hop questions across many documents

Lowering reasoning effort is also the primary way to control the number of GPT-5 tokens consumed per query — fewer planning passes and less iteration directly reduce both latency and inference cost, without needing to touch the underlying index.

Multi-hop queries in practice

The reason this matters versus classic RAG: a query like "which vendors had payment terms that changed between the Q3 and Q4 renewals" can't be answered by a single embedding lookup. With medium or high reasoning effort, the knowledge base decomposes it into sub-queries (find Q3 renewals, find Q4 renewals, compare payment terms fields), runs them in parallel against the knowledge source, and only synthesizes a final answer once it judges the retrieved context sufficient — re-querying automatically if it isn't.

References

Building a Stateful Multi-Agent System with LangGraph and Microsoft Foundry Agent Service, From Scratch

Jubin Soni — Sun, 19 Jul 2026 03:27:09 +0000

LangGraph gives you fine-grained control over agent state and routing logic in code. Microsoft Foundry gives you managed hosting, identity, tracing, and one-click publishing to Teams and M365 Copilot. Increasingly, teams want both: LangGraph's graph-based orchestration for the hard routing logic, and Foundry's operational layer for everything downstream of "it works on my machine."

This is a from-scratch walkthrough of wiring the two together — building a small multi-agent support triage system in LangGraph, then hosting it as a Foundry Hosted Agent.

Architecture

The system has three LangGraph nodes: a triage node that classifies the incoming request, and two specialist nodes (billing, technical) that the triage node routes to conditionally.

Once this graph is compiled, it doesn't run as a standalone script — it gets wrapped and exposed through Foundry's hosting layer, which handles the runtime, session management, and identity so the graph itself stays free of infrastructure concerns.

Step 1: Define the graph state and nodes

from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END
from langchain_azure_ai.chat_models import AzureAIChatCompletionsModel
import os

class TicketState(TypedDict):
    user_input: str
    ticket_type: Literal["billing", "technical", ""]
    response: str

model = AzureAIChatCompletionsModel(
    endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
    model=os.environ["FOUNDRY_MODEL_NAME"],
)

def triage_node(state: TicketState) -> TicketState:
    result = model.invoke(
        f"Classify this support ticket as 'billing' or 'technical' only: {state['user_input']}"
    )
    ticket_type = "billing" if "billing" in result.content.lower() else "technical"
    return {**state, "ticket_type": ticket_type}

def billing_node(state: TicketState) -> TicketState:
    result = model.invoke(f"You are a billing specialist. Resolve: {state['user_input']}")
    return {**state, "response": result.content}

def technical_node(state: TicketState) -> TicketState:
    result = model.invoke(f"You are a technical specialist. Resolve: {state['user_input']}")
    return {**state, "response": result.content}

Step 2: Wire the conditional routing

def route(state: TicketState) -> str:
    return state["ticket_type"]

graph = StateGraph(TicketState)
graph.add_node("triage", triage_node)
graph.add_node("billing", billing_node)
graph.add_node("technical", technical_node)
graph.set_entry_point("triage")
graph.add_conditional_edges("triage", route, {"billing": "billing", "technical": "technical"})
graph.add_edge("billing", END)
graph.add_edge("technical", END)

compiled_graph = graph.compile()

Notice the model calls point at FOUNDRY_PROJECT_ENDPOINT, not directly at Azure OpenAI or any single provider's SDK. This is what lets the same graph run locally against a Foundry project and later behind the hosted runtime without changing model-calling code.

Step 3: Expose the graph for Foundry hosting

Foundry's langchain_azure_ai.agents.hosting package wraps a compiled LangGraph graph so it can be deployed as a Hosted Agent, speaking either the Responses or Invocations protocol.

from langchain_azure_ai.agents.hosting import LangGraphResponsesHost

host = LangGraphResponsesHost(graph=compiled_graph)
app = host.create_app()

That app is what gets deployed — Foundry manages the runtime, session handling, and scaling from there. Locally you can run it with uvicorn app:app --reload and test over HTTP before deploying with azd ai agent deploy or the Foundry Toolkit extension for VS Code.

Local tools vs. built-in tools

One distinction that trips people up: LangGraph tools you attach in code run locally, colocated with your graph. Foundry's own built-in tools (file search, code interpreter, etc.) run server-side inside the Agent Service and don't add a node to your graph at all — they're invoked transparently when the service handles a request.

Aspect	Local (LangGraph) tools	Built-in Foundry tools
Execution location	Colocated with your graph process	Server-side in Foundry Agent Service
Graph representation	Adds a tool node	No node added
Applicable to	Any LangGraph/LangChain tool	Foundry agents only
Best for	Deterministic business logic, custom APIs	File search, code interpreter, managed retrieval

Tracing across two systems

Because the LangGraph portion runs outside Foundry's native agent runtime during local development, its traces don't automatically appear in the Foundry portal. You view LangGraph-level traces through Azure Monitor / Application Insights (Investigate → Agents), while Foundry-native agent traces show up directly in the Foundry portal. Once deployed as a Hosted Agent, LangChain and LangGraph applications can also be registered in the Foundry Control Plane for unified governance and trace visibility.

What this buys you in production

Once deployed, the same graph becomes available as a turnkey Application across Teams, M365 Copilot, and Bot Service channels without additional publishing code — Foundry handles the channel plumbing, while your LangGraph code stays focused purely on routing and reasoning logic.

References

Implementing Zero-Trust Networking and Identity for Microsoft Foundry Agents

Jubin Soni — Sat, 04 Jul 2026 21:38:48 +0000

Once your multi-agent system (Parts 6-8) is functionally solid, the question that comes up in every enterprise security review is the same: how do you know an agent is only doing what it's authorized to do, over a network path you actually trust, with an audit trail that holds up when someone asks "what did this agent do and why" six months later? This post covers private networking, Entra Agent ID, and audit trail design.

Zero-trust network boundary

Private endpoints: keeping traffic off the public internet

By default, calls from your application to Foundry Models travel over the public internet (encrypted, but publicly routable). For regulated workloads, route through Private Link instead:

resource foundryPrivateEndpoint 'Microsoft.Network/privateEndpoints@2023-04-01' = {
  name: 'pe-foundry-project'
  location: location
  properties: {
    subnet: {
      id: subnetId
    }
    privateLinkServiceConnections: [
      {
        name: 'foundry-connection'
        properties: {
          privateLinkServiceId: foundryResourceId
          groupIds: ['account']
        }
      }
    ]
  }
}

Pair this with a Managed VNET at the project level so that outbound calls from your compute (the orchestration layer running the agent code from Part 7) never leave the private network boundary. The failure mode to check for: a dependency library making an unexpected public call (a package doing telemetry callback, for instance) that bypasses your intended network boundary entirely — audit your dependencies' network behavior, not just your own code's.

Entra Agent ID: identity for autonomous agents, not just services

Traditional managed identity was designed for services calling APIs on a fixed schedule with predictable behavior. Entra Agent ID extends this model specifically for autonomous agents that make independent decisions — the identity model needs to answer not just "is this call authenticated" but "was this agent authorized to take this specific action in this specific context."

from azure.identity import DefaultAzureCredential

class AgentIdentityContext:
    def __init__(self, agent_id: str, allowed_actions: list[str]):
        self.agent_id = agent_id
        self.allowed_actions = allowed_actions
        self.credential = DefaultAzureCredential()

    def authorize_action(self, action: str, resource_scope: str) -> bool:
        if action not in self.allowed_actions:
            log_authorization_denial(self.agent_id, action, resource_scope)
            return False
        token = self.credential.get_token(resource_scope)
        return token is not None

def execute_agent_action(agent_context: AgentIdentityContext, action: str, resource_scope: str, fn):
    if not agent_context.authorize_action(action, resource_scope):
        raise PermissionError(f"Agent {agent_context.agent_id} not authorized for {action}")
    return fn()

The key design point: each agent in your multi-agent chain gets its own scoped identity with its own explicit allowed_actions list, rather than all agents sharing one broad service principal. This means a compromised or misbehaving refund agent can't accidentally (or maliciously, if prompt-injected) invoke actions scoped only to the fraud-check agent.

Security control	What it protects against	Enforced where
Private Link / Managed VNET	Traffic leaving the trusted network boundary	Network layer
Entra Agent ID scoped identity	An agent invoking actions outside its role	Identity/authorization layer, in code
Structured audit log with reasoning	Inability to explain why an action was taken	Application logging layer
Authorization enforced in code, not model claims	Prompt injection claiming false authorization	Application logic, never the model's own output

Data residency and customer-managed keys

Private networking and identity handle "who can reach this system and act as whom." A separate, equally common enterprise requirement is "where does this data physically live, and who holds the encryption keys" — data residency and customer-managed keys (CMK) address this and get missed by teams who've handled networking and identity but stop there.

For data residency, confirm which region your Foundry project, its underlying Azure AI Search index, and any storage backing Foundry IQ actually run in — these can silently default to a different region than your primary application if not explicitly configured, which is a real problem for workloads subject to data-sovereignty requirements (GDPR-adjacent obligations, government contracts with residency clauses, etc.):

resource foundryProject 'Microsoft.CognitiveServices/accounts/projects@2024-10-01' = {
  name: 'your-project'
  location: 'westeurope'  // must match your data residency requirement explicitly —
                            // don't rely on the default location of the parent resource group
  properties: {
    // ...
  }
}

For customer-managed keys, the default is Microsoft-managed encryption at rest, which is adequate for most workloads but insufficient for some regulated industries that require the ability to revoke access to data by rotating or deleting a key the organization itself controls:

resource foundryAccount 'Microsoft.CognitiveServices/accounts@2024-10-01' = {
  properties: {
    encryption: {
      keySource: 'Microsoft.KeyVault'
      keyVaultProperties: {
        keyVaultUri: keyVaultUri
        keyName: 'foundry-cmk'
      }
    }
  }
}

Neither of these is something you retrofit easily after data has already been written under the default configuration — data residency and CMK are decisions to make explicitly during initial provisioning of the Foundry project, not something to leave as "we'll configure that before the security review." If your compliance requirements are still being finalized when you provision, default to the more restrictive configuration (explicit region pinning, CMK) rather than the platform default, since loosening a restriction later is far easier than migrating already-written data into a stricter configuration retroactively.

Audit trails that actually hold up

A log line saying "agent called refund API" is not an audit trail — when someone asks "why did the agent decide to issue this refund," you need the reasoning captured, not just the action:

import json
import time

def log_agent_decision(agent_id, state, decision, reasoning, authorized_by):
    audit_entry = {
        "timestamp": time.time(),
        "agent_id": agent_id,
        "task_id": state.task_id,
        "decision": decision,
        "reasoning_summary": reasoning,   # the model's stated rationale, captured verbatim
        "state_snapshot": {
            "order_id": state.order_id,
            "fraud_check_result": state.fraud_check_result,
        },
        "authorized_by": authorized_by,   # which identity/policy allowed this
        "input_hash": hash_conversation_context(state),  # for reproducibility without storing PII long-term
    }
    write_to_audit_log(audit_entry)

Three things a real audit trail needs that a simple action log doesn't: the model's stated reasoning at the decision point (not just the outcome), the authorization context (which identity/policy permitted the action), and enough state snapshot to reconstruct why the decision made sense given the inputs — without necessarily storing full raw PII long-term, which is why hashing the input context is often the right tradeoff between auditability and data minimization.

Testing the security boundary, not just the happy path

Write tests that deliberately attempt to violate the boundaries you've set up:

Attempt to call an action outside an agent's allowed_actions list and confirm it's denied and logged.
Attempt a request that would resolve to a public endpoint and confirm the network boundary rejects it.
Simulate a prompt-injection attempt that tries to get an agent to claim authorization for an action it doesn't have, and confirm the authorization check (not the model's own claim) is what gates execution.

That last point matters most: authorization must be enforced in code against the identity's actual permitted action list, never based on what the model says about its own permissions in its output text.

References

Private Link for Azure AI Foundry: https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/configure-private-link
Managed VNET for Foundry projects: https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/configure-managed-network
Microsoft Entra Agent ID: https://learn.microsoft.com/en-us/entra/identity/agent-identity/overview
Zero Trust architecture guidance: https://learn.microsoft.com/en-us/security/zero-trust/zero-trust-overview

Building Evaluation, Cost Governance, and Observability for a Multi-Agent System in Microsoft Foundry

Jubin Soni — Sat, 04 Jul 2026 21:33:37 +0000

This closes out the series' capstone: the multi-agent customer support system built across Parts 6-9, now hardened with evaluation, cost governance, and observability so it can actually run in production with an on-call rotation behind it, not just in a demo environment.

Continuous evaluation pipeline

Evaluation: measuring quality continuously, not just at launch

A one-time eval before launch tells you nothing about drift once real traffic — and real edge cases — start hitting the system. Set up a continuous evaluation pipeline using a G-Eval-style approach, where a separate model scores production outputs against explicit criteria:

eval_criteria = {
    "correctness": "Does the response accurately reflect the order/refund status retrieved from the tools?",
    "escalation_appropriateness": "If the case was ambiguous or high-risk, did the agent escalate to a human rather than resolving it alone?",
    "tone": "Is the response professional and appropriately empathetic given the customer's stated frustration level?",
}

def geval_score(response, context, criterion_name, criterion_description, eval_model_client):
    prompt = f"""Evaluate the following response against this criterion: {criterion_description}
Context: {context}
Response: {response}
Score from 1-5 and give one sentence of reasoning. Return JSON: {{"score": int, "reasoning": str}}"""
    result = eval_model_client.complete(prompt)
    return json.loads(result)

def run_continuous_eval(sample_of_production_traffic):
    scores = {crit: [] for crit in eval_criteria}
    for interaction in sample_of_production_traffic:
        for crit_name, crit_desc in eval_criteria.items():
            result = geval_score(interaction.response, interaction.context, crit_name, crit_desc, eval_model_client)
            scores[crit_name].append(result["score"])
    return {crit: sum(vals) / len(vals) for crit, vals in scores.items()}

Sample a percentage of real production traffic daily (not just synthetic test cases) and track these scores over time. A drop in escalation_appropriateness specifically is the metric most worth alerting on — it's a direct proxy for the system doing something risky without a human check, which is exactly the failure mode the recovery and authorization work in Parts 7 and 9 was designed to prevent.

Cost governance: PTU vs. pay-as-you-go, decided with real math

For a system with predictable, sustained traffic (which a production support system should have), Provisioned Throughput (PTU) usually beats pay-as-you-go on cost — but the crossover point depends on your actual volume:

def compare_ptu_vs_payg(monthly_token_volume, ptu_monthly_cost, payg_per_1k_tokens):
    payg_monthly_cost = (monthly_token_volume / 1000) * payg_per_1k_tokens
    return {
        "payg_monthly": payg_monthly_cost,
        "ptu_monthly": ptu_monthly_cost,
        "recommendation": "ptu" if ptu_monthly_cost < payg_monthly_cost else "payg",
        "breakeven_tokens": (ptu_monthly_cost / payg_per_1k_tokens) * 1000,
    }

Run this quarterly, not once — traffic volume for a maturing production system tends to grow, and the PTU crossover point is usually reached faster than teams expect once an agent system is handling a meaningful fraction of real support volume.

Chargeback tagging: attributing cost to the right owner

With multiple agents (fraud-check, refund, notification) potentially running on shared compute, tag at the project level so cost attribution doesn't require manual reconciliation later:

resource_tags = {
    "business-unit": "customer-support",
    "system": "multi-agent-refund-flow",
    "environment": "production",
    "cost-center": "CC-4471",
}

Apply these consistently at the Azure resource level (not just in application logs) so Cost Management reports can be filtered directly without a separate reconciliation step — this is the difference between a chargeback model that's usable monthly versus one that requires a data-engineering project every quarter.

Dashboard signal	Source	What it indicates
Request-level trace	Part 2 tracing patterns	Latency and failure location per agent step
Authorization denials	Part 9 identity logging	Potential security issue, not just a bug
Escalation rate vs. appropriateness score	Eval pipeline + agent logs	Whether the system is escalating correctly
Cost burn rate	Azure Cost Management tags	Budget overage risk before month-end

Observability: the on-call-ready dashboard

Pull together the tracing work from Part 2, the authorization logging from Part 9, and the eval scores above into a single dashboard an on-call engineer can actually use at 2am:

Request-level trace: which agents were invoked, in what order, with what latency per step (from Part 2's tracing patterns).
Authorization denials: any agent attempting an action outside its scope (from Part 9) — a spike here is a security signal, not just a bug signal.
Escalation rate: percentage of interactions escalated to a human, tracked against the eval-measured escalation_appropriateness score — a rising escalation rate paired with a falling appropriateness score means the system is escalating things it shouldn't, which is its own kind of problem.
Cost burn rate: token consumption against the PTU/PAYG budget, with an alert threshold before month-end overage becomes a surprise.

A concrete incident: what the on-call runbook actually looks like

All the observability infrastructure above is only as good as the runbook someone follows at 2am when an alert fires. Here's a worked example tying every prior post together into one incident response flow, using a realistic trigger: the escalation-rate alert from the dashboard fires, showing escalations up 3x over baseline in the last 30 minutes.

Step 1 — check the authorization denial log (Part 9). A spike in escalations correlated with a spike in authorization denials usually means an agent is attempting actions outside its scope — possibly a misconfigured deployment, possibly a prompt-injection attempt. This is checked first because it's the highest-severity possible cause.

Step 2 — check the circuit breaker state (Part 7). If a downstream dependency (the fraud-check API, say) is degraded, the circuit breaker should already be routing to human escalation rather than retrying — confirm it's open and working as designed, not that agents are timing out repeatedly without the breaker engaging.

Step 3 — check the eval scores for escalation_appropriateness (this post). If the score is stable and escalations are simply more frequent, this may be a legitimate traffic pattern (a genuinely higher-risk cohort of requests, e.g., during a known incident like a payment processor outage) rather than a system problem. If the score is dropping alongside the escalation spike, the system's judgment about when to escalate may itself be degrading — this points back toward Part 5's schema validation and Part 7's handoff logic as places to check for a recent regression.

Step 4 — check recent deployments against the canary process (Part 2). Cross-reference the timestamp of the spike against any recent flow, model version, or schema change. If a change went out in the last few hours without full canary ramp-up, that's the most likely single cause, and rollback is usually faster than root-causing forward.

def incident_triage(alert_context):
    checks = [
        ("authorization_denials", check_authorization_spike),
        ("circuit_breaker_state", check_circuit_breaker_status),
        ("eval_score_trend", check_escalation_appropriateness_trend),
        ("recent_deployments", check_recent_flow_changes),
    ]
    findings = {}
    for name, check_fn in checks:
        findings[name] = check_fn(alert_context)
        if findings[name].get("severity") == "critical":
            return {"triage_result": name, "findings": findings, "action": "immediate_rollback_or_escalation"}
    return {"triage_result": "inconclusive", "findings": findings, "action": "manual_investigation"}

Writing this ordering down explicitly — check security signals before assuming it's a quality regression, check for a bad deploy before deep root-causing — is what turns nine posts' worth of individually reasonable safeguards into something an on-call engineer who didn't build the system can actually execute under pressure.

References

Azure AI Foundry evaluation SDK: https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/develop/evaluate-sdk
G-Eval and LLM-as-judge evaluation approach: https://learn.microsoft.com/en-us/azure/ai-foundry/concepts/observability
Provisioned Throughput Units (PTU) for Azure OpenAI: https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/provisioned-throughput
Azure Cost Management and tagging: https://learn.microsoft.com/en-us/azure/cost-management-billing/costs/cost-mgt-best-practices

Azure Databricks vs Microsoft Fabric: An Honest Guide to When to Use What

Jubin Soni — Mon, 29 Jun 2026 18:18:21 +0000

If you're building a data platform on Azure in 2026, you're going to be asked this question: Azure Databricks or Microsoft Fabric? Both run on Delta Lake, both integrate with ADLS Gen2, both have Spark, and both promise to be your unified data platform. The overlap is real and the marketing doesn't help.

This post is an honest breakdown of where each genuinely excels, where they overlap, and how to decide without getting lost in feature comparison tables.

Architecture Comparison

Decision Flow

Detailed Capability Comparison

Capability	Azure Databricks	Microsoft Fabric	Winner
Spark engine	Full Spark, Photon, tunable	Spark via Notebooks, less tunable	Databricks
Delta Lake	Native, full control	Via OneLake (Delta Parquet)	Tie
MLflow / MLOps	Native, full MLflow stack	Basic experiment tracking	Databricks
Model serving	Databricks Model Serving	Azure ML integration	Databricks
Power BI integration	DirectQuery via SQL Warehouse	Direct Lake (zero-copy, faster)	Fabric
SQL analytics	Serverless SQL Warehouse + Photon	SQL Analytics Endpoint	Tie
Data pipelines	Delta Live Tables, Workflows	Data Factory pipelines (mature)	Tie
Real-time intelligence	Spark Streaming + Kafka	Eventstream + KQL Database	Fabric
Setup complexity	Medium-high	Low (SaaS)	Fabric
Fine-grained governance	Unity Catalog (mature)	Purview integration (growing)	Databricks
Cost model	DBU + VM	Fabric capacity units	Comparable
Open format portability	High (standard Delta/Parquet)	Medium (OneLake but some lock-in)	Databricks

Step 1 — Reading Data from Fabric OneLake in Azure Databricks

The good news: Fabric and Databricks can share data via OneLake, which speaks Delta format. You don't have to pick one and abandon the other.

# Azure Databricks reading from Microsoft Fabric OneLake
# OneLake exposes an ABFS-compatible endpoint

# Authenticate using the workspace's Managed Identity or Service Principal
tenant_id     = dbutils.secrets.get("kv-scope", "sp-tenant-id")
client_id     = dbutils.secrets.get("kv-scope", "sp-client-id")
client_secret = dbutils.secrets.get("kv-scope", "sp-client-secret")

# OneLake uses the same ABFS protocol as ADLS Gen2
fabric_workspace_id = "your-fabric-workspace-guid"
lakehouse_name      = "your-lakehouse-name"
onelake_host        = "onelake.dfs.fabric.microsoft.com"

spark.conf.set(f"fs.azure.account.auth.type.{onelake_host}",             "OAuth")
spark.conf.set(f"fs.azure.account.oauth.provider.type.{onelake_host}",
               "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider")
spark.conf.set(f"fs.azure.account.oauth2.client.id.{onelake_host}",      client_id)
spark.conf.set(f"fs.azure.account.oauth2.client.secret.{onelake_host}",  client_secret)
spark.conf.set(f"fs.azure.account.oauth2.client.endpoint.{onelake_host}",
               f"https://login.microsoftonline.com/{tenant_id}/oauth2/token")

# Read a Delta table from Fabric Lakehouse
fabric_path = f"abfss://{fabric_workspace_id}@{onelake_host}/{lakehouse_name}.Lakehouse/Tables/sales_gold"

fabric_df = spark.read.format("delta").load(fabric_path)
print(f"Rows from Fabric Lakehouse: {fabric_df.count()}")
fabric_df.show(5)

Step 2 — Writing Databricks Results Back to OneLake

Run heavy ML feature engineering in Databricks, write results back to OneLake so Fabric Power BI can consume them via Direct Lake — zero-copy, sub-second dashboard refresh.

from pyspark.sql.functions import current_timestamp, lit

# Run your Databricks feature engineering / ML inference here
result_df = spark.table("production.gold.churn_predictions") \
    .withColumn("_computed_at", current_timestamp()) \
    .withColumn("_source",      lit("databricks-inference-job"))

# Write back to Fabric OneLake as Delta
output_path = f"abfss://{fabric_workspace_id}@{onelake_host}/{lakehouse_name}.Lakehouse/Tables/churn_predictions"

result_df.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \
    .save(output_path)

print(f"Written {result_df.count()} rows to Fabric OneLake.")
print("Power BI Direct Lake will pick up changes automatically.")

Step 3 — When to Use Fabric Notebooks vs Databricks Notebooks

Not everything needs Databricks. Fabric Notebooks are good enough for lighter data prep that feeds Power BI reports.

# This kind of transformation is fine in Fabric Notebooks
# Use Fabric when: output goes directly to Power BI, team is analytics-focused,
# no MLflow tracking needed, data volume < 100GB

# Fabric Notebook (PySpark — same syntax as Databricks)
from pyspark.sql.functions import col, sum as _sum, date_trunc

df = spark.read.format("delta").load("Tables/sales_silver")

summary = df \
    .withColumn("month", date_trunc("month", col("sale_ts"))) \
    .groupBy("month", "region", "product_category") \
    .agg(_sum("revenue").alias("monthly_revenue")) \
    .orderBy("month", "region")

# Write to Lakehouse table — Power BI picks it up via Direct Lake
summary.write.format("delta").mode("overwrite").saveAsTable("monthly_revenue_summary")

# Use Databricks when: MLflow tracking needed, complex ML pipeline,
# Unity Catalog governance required, data volume > 1TB, streaming workloads

When to Use Which: Decision Framework

# Use this as a mental checklist when deciding

DATABRICKS_STRENGTHS = [
    "Complex ML pipelines with MLflow experiment tracking",
    "Production model serving with A/B testing",
    "Fine-grained governance via Unity Catalog (row/column security)",
    "Spark Structured Streaming with Kafka / Event Hub",
    "Very large scale ETL (multi-TB, complex joins)",
    "Open-source tool integrations (dbt, Great Expectations, etc.)",
    "Multi-cloud or portability requirements",
]

FABRIC_STRENGTHS = [
    "Power BI as the primary consumption layer (Direct Lake = fastest)",
    "Analytics-focused teams without deep Spark expertise",
    "Microsoft 365 integration (Teams, SharePoint data sources)",
    "Real-time dashboards via Eventstream + KQL",
    "Fabric Data Factory for straightforward ELT pipelines",
    "Lower operational overhead — fully SaaS managed",
    "Already licensed via Microsoft 365 E5 / Fabric capacity",
]

BOTH_TOGETHER = [
    "Heavy ML/MLOps in Databricks, results published to OneLake for Power BI",
    "Fabric Data Factory for ingestion, Databricks for complex transformation",
    "Unity Catalog governing Databricks tables, Fabric consuming via shortcuts",
]

Things to Watch in Production

OneLake shortcuts are the integration bridge. Fabric Lakehouses support shortcuts that point to external Delta tables in ADLS Gen2 — the same storage Databricks writes to. This means Databricks writes once and Fabric reads without data movement. Set up shortcuts rather than copying data between platforms.

Unity Catalog doesn't govern Fabric. Your row-level security and column masks in Unity Catalog do not apply when Fabric reads the same underlying Delta files directly. If governance is critical, either run everything through Databricks or replicate governance rules in Fabric's permission model.

Fabric capacity units and Databricks DBUs are both usage-based but measure differently. Don't try to compare them directly. Run the same workload in both and compare wall-clock time and cost on your actual data sizes.

Fabric ML is improving fast but isn't MLflow. As of early 2026, Fabric ML experiment tracking is functional but doesn't have the depth of MLflow's model registry, artifact storage, or model serving. If MLOps maturity matters, stay on Databricks for ML.

Wrapping Up

The honest answer is: most mature Azure data platforms in 2026 use both. Azure Databricks for ML, complex transformations, governance, and streaming. Microsoft Fabric for Power BI-first analytics, simpler pipelines, and teams that don't need the full Databricks stack. OneLake shortcuts and the shared Delta format make them composable rather than competitive.

Pick based on your primary consumer: if it's Power BI dashboards, start with Fabric. If it's ML models and data products, start with Databricks. When you need both, they integrate cleanly.

References

MLOps: Building a CI/CD Pipeline for ML Models on Azure Databricks

Jubin Soni — Sun, 28 Jun 2026 15:03:49 +0000

Most ML teams are great at training models. Very few are great at shipping them. The gap between a notebook that works and a model that reliably serves production traffic is where most ML projects stall.

In this tutorial I'll walk through building a proper CI/CD pipeline for ML models on Azure Databricks using:

MLflow for experiment tracking, model versioning, and the model registry
Databricks Asset Bundles (DABs) for infrastructure-as-code and job deployment
Azure DevOps / GitHub Actions as the CI/CD orchestrator
Delta Lake as the feature and validation data store
Databricks Model Serving for the production endpoint

The use case is the same churn prediction model from the previous post, but this time we're focusing entirely on how it gets from a notebook to a production endpoint reliably and repeatably.

Architecture Overview

CI/CD Flow

Pipeline Stage Breakdown

Stage	Tool	What happens	Gate to next stage
CI	GitHub Actions	Lint, unit tests, bundle validate	All tests green
Train	Databricks Job	Full training run, MLflow logging	Job exits 0
Register	MLflow Registry	Model versioned and moved to Staging	Auto on train success
Validate	Databricks Job	Metric thresholds checked on holdout set	ROC-AUC >= 0.80
Promote	MLflow Registry	Model moved to Production	Validation passes
Deploy	Model Serving	Endpoint updated to new model version	Promotion complete
Monitor	Databricks Lakehouse Monitoring	Drift and accuracy tracked post-deploy	Ongoing

Step 1 — Project Structure with Databricks Asset Bundles

Databricks Asset Bundles (DABs) let you define your jobs, clusters, and pipelines as code and deploy them via CLI. This is your IaC layer.

# databricks.yml
bundle:
  name: churn-prediction

variables:
  env:
    default: dev

targets:
  dev:
    mode: development
    default: true
    workspace:
      host: https://adb-xxxx.azuredatabricks.net

  staging:
    workspace:
      host: https://adb-xxxx.azuredatabricks.net

  production:
    mode: production
    workspace:
      host: https://adb-xxxx.azuredatabricks.net

resources:
  jobs:
    training_job:
      name: churn-training-${var.env}
      tasks:
        - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/01_feature_engineering.py
          new_cluster:
            spark_version: 14.3.x-scala2.12
            node_type_id: Standard_DS3_v2
            num_workers: 4

        - task_key: train_and_register
          depends_on:
            - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/02_train_and_register.py
          new_cluster:
            spark_version: 14.3.x-scala2.12
            node_type_id: Standard_DS3_v2
            num_workers: 2

    validation_job:
      name: churn-validation-${var.env}
      tasks:
        - task_key: validate_model
          notebook_task:
            notebook_path: ./notebooks/03_validate_and_promote.py
          new_cluster:
            spark_version: 14.3.x-scala2.12
            node_type_id: Standard_DS3_v2
            num_workers: 2

Step 2 — Training Job with MLflow Autologging

Keep the training notebook clean and focused. Let MLflow autologging handle the heavy lifting on metric and param capture.

# notebooks/02_train_and_register.py
import mlflow
import mlflow.sklearn
from mlflow.models.signature import infer_signature
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, f1_score, precision_score, recall_score
from delta.tables import DeltaTable
import pandas as pd

mlflow.set_experiment('/churn-prediction/ci-cd-pipeline')
mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True)

# Read Gold features and capture Delta version for reproducibility
gold_table    = DeltaTable.forName(spark, 'churn.gold.features')
delta_version = gold_table.history(1).select('version').collect()[0][0]

features_pdf = spark.table('churn.gold.features') \
    .filter("feature_date = current_date()") \
    .toPandas()

FEATURE_COLS = [
    'total_events', 'total_sessions', 'distinct_products',
    'events_last_30d', 'events_last_90d', 'days_since_last_event',
    'customer_tenure_days', 'avg_events_per_day',
    'recency_tier', 'engagement_score',
]
TARGET = 'churned'

X = features_pdf[FEATURE_COLS]
y = features_pdf[TARGET]
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

with mlflow.start_run() as run:
    params = {
        'n_estimators':  200,
        'max_depth':     5,
        'learning_rate': 0.05,
        'subsample':     0.8,
    }
    model = GradientBoostingClassifier(**params, random_state=42)
    model.fit(X_train, y_train)

    y_prob = model.predict_proba(X_test)[:, 1]
    y_pred = model.predict(X_test)

    mlflow.log_param('delta_feature_version', delta_version)
    mlflow.log_param('git_sha', dbutils.widgets.get('git_sha'))  # passed in by CI
    mlflow.log_metric('roc_auc',   roc_auc_score(y_test, y_prob))
    mlflow.log_metric('f1_score',  f1_score(y_test, y_pred))
    mlflow.log_metric('precision', precision_score(y_test, y_pred))
    mlflow.log_metric('recall',    recall_score(y_test, y_pred))

    signature = infer_signature(X_train, y_pred)
    mlflow.sklearn.log_model(
        model,
        artifact_path='churn-gbm',
        signature=signature,
        registered_model_name='churn-prediction-gbm',
        await_registration_for=300,
    )

    # Pass run ID to downstream tasks via job output
    dbutils.jobs.taskValues.set(key='run_id', value=run.info.run_id)
    print(f"Training complete. Run ID: {run.info.run_id}")

Step 3 — Validation Job: Gate on Metrics Before Promoting

Never auto-promote based on a successful training run alone. Always validate on a held-out or recent dataset and check against a threshold before touching the Production alias.

# notebooks/03_validate_and_promote.py
import mlflow
from mlflow import MlflowClient
from sklearn.metrics import roc_auc_score
import pandas as pd

client    = MlflowClient()
MODEL_NAME = 'churn-prediction-gbm'

# Pick up run_id from the upstream training task
run_id = dbutils.jobs.taskValues.get(
    taskKey='train_and_register', key='run_id'
)

# Thresholds — fail the job if any metric is below these
THRESHOLDS = {
    'roc_auc':   0.80,
    'f1_score':  0.72,
    'precision': 0.70,
}

run     = client.get_run(run_id)
metrics = run.data.metrics

print("Validating metrics against thresholds...")
failures = []
for metric, threshold in THRESHOLDS.items():
    actual = metrics.get(metric, 0)
    status = 'PASS' if actual >= threshold else 'FAIL'
    print(f"  {metric}: {actual:.4f} (threshold: {threshold}) -> {status}")
    if actual < threshold:
        failures.append(f"{metric}={actual:.4f} below threshold {threshold}")

if failures:
    raise Exception(f"Validation failed: {', '.join(failures)}")

# Promote to Production alias if all thresholds pass
model_version = client.search_model_versions(
    filter_string=f"run_id='{run_id}'"
)[0].version

client.set_registered_model_alias(
    name=MODEL_NAME,
    alias='Production',
    version=model_version,
)

print(f"Model version {model_version} promoted to Production alias.")

Step 4 — GitHub Actions CI/CD Workflow

This is the glue that ties everything together. One workflow handles PR validation; the other handles deployment on merge to main.

# .github/workflows/mlops.yml
name: MLOps CI/CD

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]

env:
  DATABRICKS_HOST:  ${{ secrets.DATABRICKS_HOST }}
  DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}

jobs:
  ci:
    name: Lint and Test
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: pip install databricks-cli databricks-sdk pytest flake8 mlflow

      - name: Lint
        run: flake8 notebooks/ src/ --max-line-length=120

      - name: Unit tests
        run: pytest tests/unit/ -v

      - name: Validate DAB bundle
        run: databricks bundle validate --target staging

  train-and-deploy:
    name: Train, Validate, Deploy
    runs-on: ubuntu-latest
    needs: ci
    if: github.ref == 'refs/heads/main'
    steps:
      - uses: actions/checkout@v4

      - name: Install Databricks CLI
        run: pip install databricks-cli databricks-sdk

      - name: Deploy bundle to staging
        run: databricks bundle deploy --target staging

      - name: Run training job
        run: |
          JOB_RUN_ID=$(databricks bundle run training_job \
            --var "git_sha=${{ github.sha }}" \
            --output json | jq -r '.run_id')
          echo "TRAINING_RUN_ID=$JOB_RUN_ID" >> $GITHUB_ENV

      - name: Run validation job
        run: |
          databricks bundle run validation_job --target staging
          echo "Validation passed. Promoting to production."

      - name: Deploy bundle to production
        run: databricks bundle deploy --target production

      - name: Update Model Serving endpoint
        run: |
          python scripts/update_serving_endpoint.py \
            --model-name churn-prediction-gbm \
            --alias Production

Step 5 — Update the Serving Endpoint

The final step in the pipeline updates the Databricks Model Serving endpoint to point at the newly promoted Production model version.

# scripts/update_serving_endpoint.py
import argparse
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ServedModelInput, EndpointCoreConfigInput
from mlflow import MlflowClient

def update_serving_endpoint(model_name: str, alias: str):
    w      = WorkspaceClient()
    client = MlflowClient()

    # Resolve alias to concrete version
    model_version = client.get_model_version_by_alias(model_name, alias).version
    print(f"Deploying {model_name} version {model_version} (alias: {alias})")

    endpoint_name = 'churn-prediction-endpoint'
    served_model  = ServedModelInput(
        model_name=model_name,
        model_version=model_version,
        workload_size='Small',
        scale_to_zero_enabled=True,
    )

    try:
        # Update existing endpoint
        w.serving_endpoints.update_config(
            name=endpoint_name,
            served_models=[served_model],
        )
        print(f"Endpoint '{endpoint_name}' updated to version {model_version}.")
    except Exception:
        # Create if it doesn't exist yet
        w.serving_endpoints.create(
            name=endpoint_name,
            config=EndpointCoreConfigInput(served_models=[served_model]),
        )
        print(f"Endpoint '{endpoint_name}' created with version {model_version}.")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-name', required=True)
    parser.add_argument('--alias',      required=True)
    args = parser.parse_args()
    update_serving_endpoint(args.model_name, args.alias)

Tool Comparison

Tool	Role in pipeline	Why not the alternative
Databricks Asset Bundles	IaC for jobs and clusters	Terraform Databricks provider (more verbose, no native notebook support), manual UI (not reproducible)
MLflow Registry aliases	Production / Staging promotion	Stage-based promotion (deprecated in MLflow 2.x), manual version tracking (error prone)
GitHub Actions	CI/CD orchestrator	Azure DevOps (works equally well, swap yml syntax), Jenkins (more ops overhead)
Metric gate in validation job	Automated quality check	Manual review (blocks velocity), no gate at all (risky)
Databricks Model Serving	Managed REST endpoint	AKS deployment (more control, much more ops), Azure ML endpoints (extra service dependency)
dbutils.jobs.taskValues	Pass run_id between tasks	Environment variables (not available cross-task in DABs), hardcoded run lookup (fragile)

Things to Watch in Production

Pin your Databricks Runtime version. Using 14.3.x-scala2.12 in your bundle ensures every training run uses the same Spark and library versions. Floating versions (latest) cause silent library drift that breaks reproducibility.

Store secrets in Azure Key Vault, not GitHub Secrets alone. GitHub Secrets work fine for CI tokens but for long-lived service principal credentials that Databricks jobs use at runtime, back them with Azure Key Vault and reference them via Databricks secret scopes.

Set a metric baseline from your current production model. Your thresholds (ROC-AUC >= 0.80) should be relative to what the current Production model achieves on the same holdout set, not an arbitrary number. Add a step in the validation job that fetches the current Production model's metrics and gates the new model against those.

Tag every MLflow run with git SHA. Logging git_sha as a param in every training run means you can always trace a model artifact back to the exact code version that produced it. Critical for incident response.

Scale to zero on serving endpoints. For non-latency-critical models, enable scale_to_zero_enabled=True on your serving endpoint. It cuts cost dramatically for endpoints that don't receive traffic 24/7.

Wrapping Up

The pattern here is straightforward: code change triggers CI, CI triggers a training job, training job registers a model, a validation job gates on metrics, and only then does the model get promoted and deployed. Nothing manual, nothing skipped.

What makes this production-grade rather than just automated is the combination of Delta versioning for feature reproducibility, MLflow aliases for clean promotion semantics, and metric-gated promotion so a worse model can never silently replace a better one.

References

Azure Databricks for MLOps and Feature Engineering at Scale with Apache Spark, Delta Lake, and MLflow

Jubin Soni — Sun, 28 Jun 2026 01:35:55 +0000

Raw data doesn't win model competitions. Features do. And when your raw data is tens of billions of rows sitting across multiple sources, you can't afford to run pandas in a notebook and call it a day.

In this tutorial I'll walk through building a production-grade feature engineering pipeline on Azure Databricks using:

Apache Spark for distributed transformation at scale
Delta Lake for reliable, versioned feature storage with ACID guarantees
MLflow for tracking feature pipeline runs, parameters, and the models trained on top of them

The use case is a customer churn prediction system, but the patterns apply to any ML feature pipeline.

Architecture Overview

The pipeline follows the Medallion Architecture — a layered approach where data gets progressively cleaner and more feature-ready as it moves from Bronze to Silver to Gold. MLflow sits across all three layers tracking every run.

Pipeline Flow

Layer Breakdown

Layer	Delta Table	What happens here	Typical latency
Bronze	`churn.bronze.events`	Raw ingest, no transforms, append only	Minutes
Silver	`churn.silver.customers`	Deduplication, null handling, schema enforcement	Minutes
Gold	`churn.gold.features`	Aggregations, window functions, encoding	Minutes to hours
MLflow Run	N/A	Training, metric logging, artifact storage	Hours
Registry	N/A	Versioned model store, stage promotion	On demand

Step 1 — Bronze Layer: Raw Ingest

The Bronze layer is append-only. No transforms. No business logic. Just get the data in and preserve it exactly as it arrived so you can always replay from source.

from pyspark.sql import SparkSession
from pyspark.sql.functions import current_timestamp, lit
from delta.tables import DeltaTable

spark = SparkSession.builder.getOrCreate()

# Read raw events from ADLS Gen2 / Event Hub / source of choice
raw_events = spark.read.format('json').load('abfss://raw@yourstorage.dfs.core.windows.net/events/')

# Add ingestion metadata — never mutate source columns
bronze_df = raw_events.withColumn('_ingested_at', current_timestamp()) \
                       .withColumn('_source', lit('events_api'))

# Write to Bronze Delta table — append only, no overwrites
bronze_df.write \
    .format('delta') \
    .mode('append') \
    .option('mergeSchema', 'true') \
    .saveAsTable('churn.bronze.events')

print(f"Bronze rows written: {bronze_df.count()}")

Why append-only? If your downstream pipeline produces bad features, you want to replay from Bronze without re-ingesting from source. Overwriting Bronze breaks that ability.

Step 2 — Silver Layer: Clean and Validate

Silver is where you enforce schema, handle nulls, deduplicate, and standardize. Think of it as your canonical, trusted dataset.

from pyspark.sql.functions import col, to_timestamp, when, trim, upper
from delta.tables import DeltaTable

bronze = spark.table('churn.bronze.events')

silver_df = bronze \
    .filter(col('customer_id').isNotNull()) \
    .filter(col('event_type').isNotNull()) \
    .dropDuplicates(['customer_id', 'event_id']) \
    .withColumn('event_ts',     to_timestamp(col('event_timestamp'))) \
    .withColumn('event_type',   upper(trim(col('event_type')))) \
    .withColumn('country_code', when(col('country').isNull(), lit('UNKNOWN'))
                                .otherwise(upper(col('country')))) \
    .select(
        'customer_id',
        'event_id',
        'event_type',
        'event_ts',
        'country_code',
        'product_id',
        'session_id',
        '_ingested_at',
    )

# Upsert into Silver using Delta MERGE — idempotent on re-runs
if DeltaTable.isDeltaTable(spark, 'churn.silver.customers'):
    silver_table = DeltaTable.forName(spark, 'churn.silver.customers')
    silver_table.alias('tgt').merge(
        silver_df.alias('src'),
        'tgt.customer_id = src.customer_id AND tgt.event_id = src.event_id'
    ).whenNotMatchedInsertAll().execute()
else:
    silver_df.write.format('delta').saveAsTable('churn.silver.customers')

print(f"Silver table updated. Total rows: {spark.table('churn.silver.customers').count()}")

Step 3 — Gold Layer: Feature Engineering

This is the heart of the pipeline. We compute aggregated, windowed, and encoded features that the model will actually train on.

from pyspark.sql.functions import (
    col, count, countDistinct, sum as _sum,
    avg, datediff, max as _max, min as _min,
    current_date, expr, when
)
from pyspark.sql.window import Window

silver = spark.table('churn.silver.customers')

# ------------------------------------------------------------------
# 1. Aggregate features per customer over 30 / 90 day windows
# ------------------------------------------------------------------
today = current_date()

agg_features = silver \
    .withColumn('days_since_event', datediff(today, col('event_ts'))) \
    .groupBy('customer_id') \
    .agg(
        count('event_id')                                          .alias('total_events'),
        countDistinct('session_id')                                .alias('total_sessions'),
        countDistinct('product_id')                                .alias('distinct_products'),
        _sum(when(col('days_since_event') <= 30, 1).otherwise(0)) .alias('events_last_30d'),
        _sum(when(col('days_since_event') <= 90, 1).otherwise(0)) .alias('events_last_90d'),
        _max('event_ts')                                           .alias('last_event_ts'),
        _min('event_ts')                                           .alias('first_event_ts'),
    ) \
    .withColumn('days_since_last_event', datediff(today, col('last_event_ts'))) \
    .withColumn('customer_tenure_days',  datediff(today, col('first_event_ts'))) \
    .withColumn('avg_events_per_day',
        col('total_events') / (col('customer_tenure_days') + 1))

# ------------------------------------------------------------------
# 2. Encode churn risk tier as ordinal feature
# ------------------------------------------------------------------
feature_df = agg_features \
    .withColumn('recency_tier',
        when(col('days_since_last_event') <= 7,  lit(3))   # active
       .when(col('days_since_last_event') <= 30, lit(2))   # at risk
       .otherwise(lit(1))                                   # churned
    ) \
    .withColumn('engagement_score',
        (col('events_last_30d') * 0.6 + col('events_last_90d') * 0.4) /
        (col('customer_tenure_days') + 1)
    )

# ------------------------------------------------------------------
# 3. Write to Gold feature store — overwrite with partition by date
# ------------------------------------------------------------------
feature_df \
    .withColumn('feature_date', current_date()) \
    .write \
    .format('delta') \
    .mode('overwrite') \
    .option('replaceWhere', f"feature_date = '{today}'") \
    .saveAsTable('churn.gold.features')

print(f"Gold features written: {feature_df.count()} customers")

Step 4 — MLflow: Track the Training Run

With features in Gold, we hand off to MLflow to train, track, and register the model. Notice we log the Delta table version so we can always reproduce exactly which feature snapshot trained which model.

import mlflow
import mlflow.sklearn
from mlflow.models.signature import infer_signature
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, f1_score
import pandas as pd

mlflow.set_experiment('/churn-prediction/feature-pipeline')

# Read Gold features — capture Delta version for reproducibility
gold_table  = DeltaTable.forName(spark, 'churn.gold.features')
delta_version = gold_table.history(1).select('version').collect()[0][0]

features_pdf = spark.table('churn.gold.features').toPandas()

FEATURE_COLS = [
    'total_events', 'total_sessions', 'distinct_products',
    'events_last_30d', 'events_last_90d', 'days_since_last_event',
    'customer_tenure_days', 'avg_events_per_day',
    'recency_tier', 'engagement_score',
]
TARGET = 'churned'

X = features_pdf[FEATURE_COLS]
y = features_pdf[TARGET]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

with mlflow.start_run(run_name=f'gbm-features-v{delta_version}') as run:

    params = {'n_estimators': 200, 'max_depth': 5, 'learning_rate': 0.05}
    model  = GradientBoostingClassifier(**params, random_state=42)
    model.fit(X_train, y_train)

    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1]

    # Log everything
    mlflow.log_params(params)
    mlflow.log_metric('roc_auc', roc_auc_score(y_test, y_prob))
    mlflow.log_metric('f1_score', f1_score(y_test, y_pred))
    mlflow.log_param('delta_feature_version', delta_version)
    mlflow.log_param('feature_columns', FEATURE_COLS)
    mlflow.log_param('training_rows', len(X_train))

    # Log model with signature
    signature = infer_signature(X_train, y_pred)
    mlflow.sklearn.log_model(
        model,
        artifact_path='churn-gbm',
        signature=signature,
        registered_model_name='churn-prediction-gbm',
    )

    print(f"Run ID: {run.info.run_id}")
    print(f"ROC-AUC: {roc_auc_score(y_test, y_prob):.4f}")
    print(f"Feature Delta version logged: {delta_version}")

Bonus: Delta Lake Time Travel for Feature Reproducibility

One of the best things about Delta Lake is time travel. If a model behaves unexpectedly in production, you can reload the exact feature snapshot it was trained on.

# Reload the exact feature version that trained a specific model run
import mlflow

run = mlflow.get_run('your-run-id-here')
feature_version = int(run.data.params['delta_feature_version'])

# Rehydrate that exact feature snapshot
historical_features = spark.read \
    .format('delta') \
    .option('versionAsOf', feature_version) \
    .table('churn.gold.features')

print(f"Loaded feature snapshot from Delta version {feature_version}")
print(f"Row count: {historical_features.count()}")

# You can now retrain on the exact same data to reproduce the result

Service Comparison

Tool	Role in pipeline	Why not the alternative
Apache Spark	Distributed feature computation	Pandas (single node, OOM at scale), Dask (less native Databricks integration)
Delta Lake	Feature storage with versioning	Parquet (no ACID, no time travel), Hive tables (no merge support)
MLflow Tracking	Experiment and param logging	Manual logging (not reproducible), W&B (extra cost, less native on Databricks)
MLflow Registry	Model versioning and promotion	Custom model store (more ops overhead)
Medallion Architecture	Pipeline layer separation	Flat pipelines (hard to debug, no replay capability)
Delta MERGE	Idempotent Silver upserts	Overwrite (destroys history), append (creates duplicates)

Things to Watch in Production

Shuffle partitions matter. Spark defaults to 200 shuffle partitions which is fine for small data but will bottleneck at scale. Set spark.conf.set("spark.sql.shuffle.partitions", "auto") on Databricks Runtime 10+ or tune it manually to 2-3x your core count.

Z-ordering on Gold features. If you're querying Gold by customer_id frequently, add OPTIMIZE churn.gold.features ZORDER BY (customer_id) after the write. This co-locates related data and cuts query times dramatically on large tables.

Log Delta version in every MLflow run. This is non-negotiable for reproducibility. Without it you can't prove which feature snapshot trained which model, which becomes a compliance problem in regulated industries.

Cluster autoscaling for feature jobs. Feature engineering jobs tend to have spiky resource needs (big during aggregation, small during writes). Enable autoscaling on your Databricks cluster and set a min/max node count rather than a fixed size.

Wrapping Up

The combination of Spark, Delta Lake, and MLflow on Databricks gives you a feature engineering pipeline that is reproducible (Delta time travel + MLflow param logging), scalable (Spark handles billions of rows), and auditable (every run is tracked, every feature version is stored).

The Medallion Architecture keeps the pipeline modular — you can rerun just the Gold layer if you change a feature definition without touching Bronze or Silver, and MLflow ties model performance back to the exact feature version that produced it.

References

Building an AI Agent That Responds to Real-Time Events with AWS Bedrock, Kinesis, DynamoDB, and S3

Jubin Soni — Sun, 28 Jun 2026 01:17:32 +0000

Most recommendation systems are batch jobs. They crunch last night's data, write a recommendations table, and serve it all day. That works fine until your user watches three thriller movies in a row at 9pm and your system is still recommending rom-coms because the batch hasn't run yet.

In this post I'll walk through building an agent system that reacts to streaming user behavior in real time using:

Amazon Kinesis to ingest and route events
AWS Lambda to process, enrich, and trigger reasoning
Amazon Bedrock as the reasoning and recommendation layer
DynamoDB to store user profiles and recommendation cache
S3 for raw event archiving and model artifacts

By the end you'll have an architecture where a user's recommendation set updates within seconds of their behavior changing.

Architecture Overview

The system has three layers:

Layer	Services	Responsibility
Ingest	Kinesis Data Streams, Kinesis Firehose	Capture and fan-out user events
Process & Reason	Lambda, Amazon Bedrock Agent	Enrich events, build context, invoke LLM
Store & Serve	DynamoDB, S3	Persist profiles, cache recs, store artifacts

The key design decision is keeping the hot path (Kinesis → Lambda → Bedrock → DynamoDB) fully async and the serving path (API → DynamoDB cache) completely decoupled. The user never waits for Bedrock to respond; they get the last cached recommendation set while a fresh one is already being computed in the background.

Event Flow

Here's what happens end to end when a user clicks on a product:

The app publishes a user.interaction event to Kinesis Data Streams
Kinesis fans the event out to two consumers: Lambda Processor and Kinesis Firehose
Firehose archives the raw event to S3 (cheap, durable, great for retraining later)
Lambda enriches the event with user history from DynamoDB User Profiles, then invokes the Bedrock Agent
The Bedrock Agent reasons over the enriched context (recent events + profile + item catalog embeddings from S3) and writes a fresh recommendation set to DynamoDB Rec Cache
The client app reads recommendations from the cache via a lightweight Lambda API — no Bedrock latency in the hot path

Code: Publishing Events to Kinesis

This is your app-side producer. Keep it thin — just serialize and publish. Do all enrichment downstream.

import boto3
import json
import uuid
from datetime import datetime, timezone

kinesis = boto3.client('kinesis', region_name='us-east-1')

def publish_interaction(user_id: str, item_id: str, event_type: str, metadata: dict = {}):
    """
    Publish a user interaction event to Kinesis Data Streams.
    Partition key is user_id so all events for a user land on the same shard.
    """
    event = {
        'event_id':   str(uuid.uuid4()),
        'user_id':    user_id,
        'item_id':    item_id,
        'event_type': event_type,          # 'click', 'purchase', 'dwell', 'skip'
        'timestamp':  datetime.now(timezone.utc).isoformat(),
        'metadata':   metadata,
    }

    response = kinesis.put_record(
        StreamName='user-interactions',
        Data=json.dumps(event).encode('utf-8'),
        PartitionKey=user_id,              # consistent routing per user
    )

    return response['SequenceNumber']


# Example call
publish_interaction(
    user_id='u_8821',
    item_id='prod_thriller_042',
    event_type='purchase',
    metadata={'price': 14.99, 'category': 'thriller', 'session_id': 'sess_xyz'}
)

Tip: Use user_id as the partition key so all events for a given user land on the same shard and arrive in order. This matters when Lambda is building a recency-ordered event window.

Code: Lambda Processor — Enrich and Invoke Bedrock

This is the core of the pipeline. The Lambda reads from the Kinesis stream, pulls user context from DynamoDB, and invokes the Bedrock Agent with a structured prompt.

import boto3
import json
import os
from datetime import datetime, timezone

dynamodb  = boto3.resource('dynamodb')
bedrock   = boto3.client('bedrock-agent-runtime', region_name='us-east-1')

profiles_table = dynamodb.Table(os.environ['PROFILES_TABLE'])   # DynamoDB User Profiles
rec_table      = dynamodb.Table(os.environ['REC_CACHE_TABLE'])  # DynamoDB Rec Cache

AGENT_ID      = os.environ['BEDROCK_AGENT_ID']
AGENT_ALIAS   = os.environ['BEDROCK_AGENT_ALIAS']
MAX_HISTORY   = 20  # last N events to include in context


def handler(event, context):
    for record in event['Records']:
        # Kinesis payload is base64-encoded
        payload = json.loads(record['kinesis']['data'])
        process_event(payload)


def process_event(payload: dict):
    user_id  = payload['user_id']
    item_id  = payload['item_id']
    evt_type = payload['event_type']

    # 1. Fetch user profile + recent history from DynamoDB
    response = profiles_table.get_item(Key={'user_id': user_id})
    profile  = response.get('Item', {'user_id': user_id, 'history': [], 'preferences': {}})

    # 2. Append current event and trim to window
    profile['history'].append({
        'item_id':    item_id,
        'event_type': evt_type,
        'timestamp':  payload['timestamp'],
        'metadata':   payload.get('metadata', {}),
    })
    profile['history'] = profile['history'][-MAX_HISTORY:]

    # 3. Write enriched profile back
    profiles_table.put_item(Item=profile)

    # 4. Build prompt for Bedrock Agent
    prompt = build_personalization_prompt(profile)

    # 5. Invoke Bedrock Agent
    agent_response = bedrock.invoke_agent(
        agentId=AGENT_ID,
        agentAliasId=AGENT_ALIAS,
        sessionId=user_id,           # session per user keeps conversational context
        inputText=prompt,
    )

    # 6. Parse streaming response chunks
    recommendations = parse_agent_response(agent_response)

    # 7. Write to recommendation cache
    rec_table.put_item(Item={
        'user_id':         user_id,
        'recommendations': recommendations,
        'generated_at':    datetime.now(timezone.utc).isoformat(),
        'ttl':             int(datetime.now(timezone.utc).timestamp()) + 3600,  # 1hr TTL
    })


def build_personalization_prompt(profile: dict) -> str:
    history_summary = '\n'.join([
        f"- [{e['event_type'].upper()}] item={e['item_id']} category={e['metadata'].get('category','unknown')}"
        for e in profile['history'][-10:]
    ])
    return f"""You are a real-time personalization agent.

User profile: {json.dumps(profile.get('preferences', {}))}

Recent interactions (most recent last):
{history_summary}

Based on this behavior, return exactly 5 personalized item recommendations as a JSON array.
Each item must include: item_id, category, reasoning (1 sentence), confidence_score (0-1).
Return only valid JSON. No explanation outside the JSON block."""


def parse_agent_response(agent_response) -> list:
    full_text = ''
    for event in agent_response['completion']:
        if 'chunk' in event:
            full_text += event['chunk']['bytes'].decode('utf-8')
    try:
        # Extract JSON from response
        start = full_text.index('[')
        end   = full_text.rindex(']') + 1
        return json.loads(full_text[start:end])
    except (ValueError, json.JSONDecodeError):
        return []

Code: Serving Recommendations via Lambda API

The serving layer never touches Bedrock. It reads purely from the DynamoDB cache, keeping p99 latency well under 10ms.

import boto3
import json
import os
from datetime import datetime, timezone

dynamodb  = boto3.resource('dynamodb')
rec_table = dynamodb.Table(os.environ['REC_CACHE_TABLE'])

FALLBACK_RECS = ['popular_001', 'popular_002', 'popular_003']  # cold-start fallback


def handler(event, context):
    user_id = event['pathParameters']['userId']

    response = rec_table.get_item(Key={'user_id': user_id})
    item     = response.get('Item')

    if not item:
        # Cold start: user has no history yet
        return api_response(200, {
            'user_id':         user_id,
            'recommendations': FALLBACK_RECS,
            'source':          'fallback',
            'generated_at':    None,
        })

    age_seconds = (
        datetime.now(timezone.utc) -
        datetime.fromisoformat(item['generated_at'])
    ).total_seconds()

    return api_response(200, {
        'user_id':         user_id,
        'recommendations': item['recommendations'],
        'source':          'cache',
        'generated_at':    item['generated_at'],
        'cache_age_sec':   int(age_seconds),
    })


def api_response(status: int, body: dict) -> dict:
    return {
        'statusCode': status,
        'headers': {
            'Content-Type':                'application/json',
            'Access-Control-Allow-Origin': '*',
        },
        'body': json.dumps(body),
    }

Service Comparison: Why Each AWS Service?

Service	Why it's here	Alternative considered
Kinesis Data Streams	Ordered, replayable, millisecond latency fan-out	SQS (no ordering guarantee per user), EventBridge (higher latency)
Kinesis Firehose	Zero-code delivery to S3 for archiving	Writing to S3 directly in Lambda (adds failure surface)
Lambda	Event-driven, scales to 0, tight Kinesis integration	ECS Fargate (overkill for stateless enrichment)
Amazon Bedrock	Managed LLM with agent runtime, no infra to maintain	Self-hosted model on SageMaker (more control, much more ops)
DynamoDB	Single-digit ms reads, TTL support, scales automatically	RDS (too slow for hot path), ElastiCache (extra cost for separate store)
S3	Cheap durable archive + model artifact store	DynamoDB for raw events (expensive and unnecessary)

Things to Watch in Production

Bedrock latency is variable. Claude Sonnet typically responds in 1-4 seconds but can spike. Since recs are written async to cache, this doesn't affect user-facing latency, but it does affect freshness. Monitor bedrock:InvokeAgent duration in CloudWatch.

Kinesis shard scaling. One shard handles 1MB/s write or 1000 records/s. At 10k active users you'll need to plan shard count carefully. Use Enhanced Fan-Out if you have multiple Lambda consumers reading the same stream.

DynamoDB TTL for cache eviction. The serving Lambda sets a 1-hour TTL on each rec entry. If Bedrock hasn't updated the cache in over an hour (e.g. Lambda errors), users fall back to the popular items list. Adjust TTL based on how stale you can tolerate.

Cold start users. New users have no history so the Bedrock prompt has nothing useful to reason over. I recommend a popularity-based fallback as shown in the serving Lambda, and switching to personalized recs after the user's first 3-5 interactions.

Wrapping Up

The pattern here is worth generalizing: keep the reasoning layer (Bedrock) fully off the hot serving path. Write results to a fast cache (DynamoDB), serve from the cache, and let the agent pipeline update it continuously in the background. This gives you the intelligence of an LLM-powered agent without the latency of one.

The same pattern applies to fraud scoring, content moderation queues, ops alerting — anywhere you need a reasoning system that reacts to real-time streams without blocking the user experience.

References

Real-Time AI Feature Engineering with Spark Structured Streaming and Databricks Feature Store

Jubin Soni — Wed, 24 Jun 2026 09:41:21 +0000

Building point-in-time correct, production-grade feature pipelines — from raw Kafka events to online feature serving in milliseconds, using Spark Structured Streaming and the Databricks Feature Store.

The Feature Engineering Problem
Architecture Overview
Feature Store Concepts: ERD
Environment Setup
Streaming Feature Pipeline
Point-in-Time Correct Training Dataset Generation
Writing Features to the Online Store
Serving Features at Inference Time
Feature Table Reference
References

The Feature Engineering Problem

Feature engineering is where most ML projects silently fail in production. Not because the model is wrong — but because the features the model sees at training time are different from the features it sees at inference time. This is called training-serving skew, and it's the #1 silent killer of ML systems.

Three specific failure modes cause it:

Online/offline inconsistency — the batch pipeline that computes training features uses different logic than the real-time service that computes inference features
Data leakage — training features accidentally include information from the future (e.g. joining on a label that was created after the event)
Feature staleness — a model trained on 30-day rolling averages is served features that are 6 hours stale because the pipeline backfills are slow

The Databricks Feature Store — now part of Unity Catalog as Feature Engineering in Unity Catalog — solves all three by:

Storing feature computation logic alongside the data (no drift between training and serving)
Enforcing point-in-time lookups during training dataset creation
Providing a unified API for both batch offline reads and low-latency online reads

Architecture Overview

Feature Store Concepts: ERD

Understanding the data model behind the Feature Store is essential for designing correct pipelines. Here's how the entities relate:

The critical relationship: a Model Version is bound to a Training Set, which records exactly which feature tables and which point-in-time lookups were used. This is how Databricks guarantees reproducibility — you can always re-create the exact training data that produced any model version.

Environment Setup

# Databricks Runtime ML 13.x+ recommended
# Feature Engineering in Unity Catalog (formerly Feature Store)

%pip install databricks-feature-engineering==0.6.0 --quiet
dbutils.library.restartPython()

from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup
from databricks.feature_engineering.entities.feature_serving_endpoint import (
    ServedEntity, EndpointCoreConfig
)
from pyspark.sql import functions as F, SparkSession
from pyspark.sql.types import (
    StructType, StructField, StringType, LongType,
    DoubleType, TimestampType, ArrayType
)
import mlflow

spark = SparkSession.builder.getOrCreate()
fe = FeatureEngineeringClient()

# Unity Catalog paths
CATALOG       = "prod"
FEATURE_DB    = f"{CATALOG}.feature_store"
EVENTS_TABLE  = f"{CATALOG}.silver.events_clean"
KAFKA_BROKER  = "kafka-broker.internal:9092"
KAFKA_TOPIC   = "user-events"

# Checkpoint locations (ADLS / S3 / GCS)
CHECKPOINT_BASE = "abfss://checkpoints@storage.dfs.core.windows.net/features"

Streaming Feature Pipeline

The streaming pipeline reads from Kafka, computes windowed aggregations using Spark's stateful streaming engine, and writes features to the Feature Store via foreachBatch. This keeps the feature table continuously fresh.

# ── Streaming Feature Pipeline ────────────────────────────────────────────────

# Step 1: Define the raw event schema from Kafka
event_schema = StructType([
    StructField("user_id",       StringType(),    False),
    StructField("event_type",    StringType(),    True),
    StructField("product_id",    StringType(),    True),
    StructField("revenue",       DoubleType(),    True),
    StructField("session_id",    StringType(),    True),
    StructField("platform",      StringType(),    True),
    StructField("event_ts",      TimestampType(), False),
])

# Step 2: Read from Kafka
raw_stream = (
    spark.readStream
        .format("kafka")
        .option("kafka.bootstrap.servers", KAFKA_BROKER)
        .option("subscribe", KAFKA_TOPIC)
        .option("startingOffsets", "latest")
        .option("failOnDataLoss", "false")
        .load()
        .select(
            F.from_json(F.col("value").cast("string"), event_schema).alias("data"),
            F.col("timestamp").alias("kafka_ts")
        )
        .select("data.*", "kafka_ts")
)

# Step 3: Apply watermark and compute windowed features
# Watermark: tolerate up to 10 minutes of late data
windowed_features = (
    raw_stream
        .withWatermark("event_ts", "10 minutes")
        .groupBy(
            F.col("user_id"),
            F.window(F.col("event_ts"), "1 hour", "15 minutes").alias("window")
        )
        .agg(
            F.count("*").alias("event_count_1h"),
            F.sum(F.when(F.col("event_type") == "purchase", F.col("revenue"))
                  .otherwise(0)).alias("revenue_1h"),
            F.countDistinct("session_id").alias("session_count_1h"),
            F.countDistinct("product_id").alias("unique_products_1h"),
            F.sum(F.when(F.col("event_type") == "purchase", 1)
                  .otherwise(0)).alias("purchase_count_1h"),
            F.first("platform").alias("last_platform"),
        )
        # Flatten window struct to scalar columns
        .withColumn("window_start", F.col("window.start"))
        .withColumn("window_end",   F.col("window.end"))
        .withColumn("feature_ts",   F.col("window.end"))   # timestamp key for PIT lookup
        .drop("window")
        # Derived features
        .withColumn("conversion_rate_1h",
            F.when(F.col("event_count_1h") > 0,
                   F.col("purchase_count_1h") / F.col("event_count_1h"))
            .otherwise(0.0))
        .withColumn("avg_revenue_per_purchase_1h",
            F.when(F.col("purchase_count_1h") > 0,
                   F.col("revenue_1h") / F.col("purchase_count_1h"))
            .otherwise(0.0))
)


# Step 4: Write to Feature Store via foreachBatch
# foreachBatch gives us transactional writes per micro-batch
def write_to_feature_store(batch_df, batch_id):
    """
    Called on each micro-batch. Merges feature data into the Feature Store
    table using merge_on keys (user_id + feature_ts).
    """
    if batch_df.isEmpty():
        return

    fe.write_table(
        name=f"{FEATURE_DB}.user_activity_features",
        df=batch_df,
        mode="merge",             # upsert: update existing, insert new
    )
    print(f"Batch {batch_id}: wrote {batch_df.count()} feature rows")


# Step 5: Create the feature table (idempotent — safe to re-run)
try:
    fe.create_table(
        name=f"{FEATURE_DB}.user_activity_features",
        primary_keys=["user_id"],
        timestamp_keys=["feature_ts"],
        schema=windowed_features.schema,
        description=(
            "Real-time user activity features computed from event stream. "
            "1-hour sliding window, refreshed every 15 minutes. "
            "Primary key: user_id. Timestamp key: feature_ts (window end)."
        ),
    )
    print("Feature table created.")
except Exception:
    print("Feature table already exists — continuing.")


# Step 6: Launch the streaming query
streaming_query = (
    windowed_features.writeStream
        .outputMode("update")               # update mode for stateful aggregations
        .option("checkpointLocation", f"{CHECKPOINT_BASE}/user_activity")
        .trigger(processingTime="5 minutes") # micro-batch every 5 min
        .foreachBatch(write_to_feature_store)
        .start()
)

print(f"Streaming query '{streaming_query.name}' running...")
print(f"Status: {streaming_query.status}")

Point-in-Time Correct Training Dataset Generation

This is the most critical part of the Feature Store. When creating training data, we must join labels to features at the timestamp of the label event — not the current time. This prevents data leakage.

# ── Point-in-Time Correct Training Dataset ────────────────────────────────────

# Step 1: Load the label dataset
# Each row = one prediction target event, with the exact timestamp
# at which a model would have needed to make a prediction.

labels_df = (
    spark.table(f"{CATALOG}.gold.churn_labels")
        .select(
            "user_id",
            "churn_label",                        # 0 = retained, 1 = churned
            F.col("observation_ts").alias("event_timestamp"),  # point-in-time anchor
            "experiment_split"                    # train/val/test
        )
        .filter(F.col("observation_ts") >= "2024-01-01")
)

print(f"Label rows: {labels_df.count():,}")
labels_df.show(5)
# +----------+-----------+---------------------+-----------------+
# | user_id  |churn_label| event_timestamp     | experiment_split|
# +----------+-----------+---------------------+-----------------+
# | u_123456 | 0         | 2024-03-15 14:22:00 | train           |
# | u_789012 | 1         | 2024-03-15 18:45:00 | train           |


# Step 2: Define feature lookups
# as_of_timestamp=None → use the label's event_timestamp (point-in-time)
# Databricks will join each label row to the feature values
# that were valid at event_timestamp — not the latest values.

feature_lookups = [
    # User activity features — 1h window features from the streaming pipeline
    FeatureLookup(
        table_name=f"{FEATURE_DB}.user_activity_features",
        feature_names=[
            "event_count_1h",
            "revenue_1h",
            "session_count_1h",
            "unique_products_1h",
            "purchase_count_1h",
            "conversion_rate_1h",
            "avg_revenue_per_purchase_1h",
            "last_platform",
        ],
        lookup_key="user_id",
        timestamp_lookup_key="event_timestamp",    # ← PIT anchor
    ),

    # User profile features — slower-changing, from batch pipeline
    FeatureLookup(
        table_name=f"{FEATURE_DB}.user_profile_features",
        feature_names=[
            "account_age_days",
            "lifetime_revenue",
            "preferred_category",
            "subscription_tier",
        ],
        lookup_key="user_id",
        timestamp_lookup_key="event_timestamp",    # ← PIT anchor
    ),

    # Transaction aggregates — 30d and 90d rolling windows
    FeatureLookup(
        table_name=f"{FEATURE_DB}.transaction_features",
        feature_names=[
            "purchase_count_30d",
            "purchase_count_90d",
            "avg_order_value_30d",
            "days_since_last_purchase",
            "category_diversity_score",
        ],
        lookup_key="user_id",
        timestamp_lookup_key="event_timestamp",
    ),
]


# Step 3: Create training dataset (Feature Store handles the PIT join)
training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=feature_lookups,
    label="churn_label",
    exclude_columns=["observation_ts", "experiment_split"],
)

# The returned DataFrame has features + labels, PIT-correct
training_df = training_set.load_df()
print(f"Training rows: {training_df.count():,}")
print(f"Training cols: {len(training_df.columns)}")
training_df.show(3)


# Step 4: Train model and log via Feature Store (preserves lineage!)
from sklearn.ensemble import GradientBoostingClassifier
import pandas as pd

train_pdf = (
    training_df
        .filter(F.col("experiment_split") == "train")
        .drop("experiment_split", "user_id")
        .fillna(0)
        .toPandas()
)

X_train = train_pdf.drop(columns=["churn_label"])
y_train = train_pdf["churn_label"]

model = GradientBoostingClassifier(
    n_estimators=300,
    learning_rate=0.05,
    max_depth=5,
    subsample=0.8,
    random_state=42,
)

with mlflow.start_run(run_name="churn-gbm-v1") as run:
    model.fit(X_train, y_train)

    # Log model via Feature Store — this records the feature lineage
    fe.log_model(
        model=model,
        artifact_path="churn_model",
        flavor=mlflow.sklearn,
        training_set=training_set,      # ← binds model to its feature lookups
        registered_model_name=f"{CATALOG}.ml.user_churn_model",
    )
    print(f"Logged model with feature lineage. Run: {run.info.run_id}")

Writing Features to the Online Store

For real-time inference, the model needs features in milliseconds — not the seconds it takes to query Delta Lake. Databricks Feature Store can publish features to an online store (DynamoDB, Cosmos DB, MySQL, etc.) for low-latency reads.

# ── Publish Features to Online Store ─────────────────────────────────────────
# Online stores are configured per feature table.
# Here we publish user_activity_features to DynamoDB for <5ms lookups.

from databricks.feature_engineering.entities.feature_store_online_table import (
    OnlineTable, OnlineTableSpec, TriggeredSchedulingPolicy
)

# Create an online table spec (backed by a serverless real-time compute layer)
online_table_spec = OnlineTableSpec(
    primary_key_columns=["user_id"],
    source_table_full_name=f"{FEATURE_DB}.user_activity_features",
    run_triggered=OnlineTableSpec.TriggeredSchedulingPolicy(),  # sync on-demand
    # OR for continuous sync:
    # run_continuous=OnlineTableSpec.ContinuousSchedulingPolicy()
)

# Create the online table (idempotent)
online_table = fe.create_online_table(spec=online_table_spec)
print(f"Online table: {online_table.name}")
print(f"Status:       {online_table.status.detailed_state}")

# Trigger an initial sync from the offline Delta table to the online store
fe.refresh_online_table(name=f"{FEATURE_DB}.user_activity_features")

Serving Features at Inference Time

At inference time, the Feature Store SDK performs automatic feature lookups, joining the incoming request data with features from the online store before passing them to the model.

# ── Real-Time Feature Serving at Inference ────────────────────────────────────

import requests, json

WORKSPACE_URL = "https://<workspace>.azuredatabricks.net"
TOKEN = dbutils.secrets.get("prod-scope", "databricks-token")


# Option 1: Model Serving with automatic feature lookup
# When you logged the model with fe.log_model(), Databricks knows which
# features to fetch. You only send the lookup key (user_id) at inference time.

def predict_churn(user_ids: list) -> list:
    """
    Send only user_id — the serving endpoint fetches features automatically
    from the online store and runs inference.
    """
    payload = {
        "dataframe_records": [
            {"user_id": uid} for uid in user_ids
        ]
    }
    resp = requests.post(
        f"{WORKSPACE_URL}/serving-endpoints/churn-predictor/invocations",
        headers={
            "Authorization": f"Bearer {TOKEN}",
            "Content-Type":  "application/json",
        },
        data=json.dumps(payload),
        timeout=5,
    )
    resp.raise_for_status()
    return resp.json()["predictions"]


# Example usage
predictions = predict_churn(["u_123456", "u_789012", "u_345678"])
for uid, pred in zip(["u_123456", "u_789012", "u_345678"], predictions):
    print(f"{uid}: churn_probability = {pred:.4f}")
# u_123456: churn_probability = 0.0821
# u_789012: churn_probability = 0.7643
# u_345678: churn_probability = 0.1209


# Option 2: Direct feature lookup via the Feature Serving endpoint
# Useful when you want raw features without running inference
def get_features(user_ids: list) -> dict:
    payload = {
        "dataframe_records": [{"user_id": uid} for uid in user_ids]
    }
    resp = requests.post(
        f"{WORKSPACE_URL}/serving-endpoints/user-features-serving/invocations",
        headers={
            "Authorization": f"Bearer {TOKEN}",
            "Content-Type":  "application/json",
        },
        data=json.dumps(payload),
        timeout=5,
    )
    return resp.json()


# Option 3: Batch scoring (offline) — uses Delta offline store
# No online store needed; reads directly from the feature table with PIT lookup
batch_labels = spark.table(f"{CATALOG}.gold.users_to_score_today") \
    .select("user_id", F.current_timestamp().alias("event_timestamp"))

batch_predictions = fe.score_batch(
    model_uri=f"models:/{CATALOG}.ml.user_churn_model@champion",
    df=batch_labels,
    result_type="double",
)

batch_predictions.select("user_id", "prediction") \
    .write.format("delta").mode("overwrite") \
    .saveAsTable(f"{CATALOG}.gold.churn_scores_daily")

Feature Table Reference

A summary of the feature tables in our pipeline, their update cadence, and their role in the ML lifecycle:

Feature Table	Primary Key	Timestamp Key	Update Method	Latency	Used In
`user_activity_features`	`user_id`	`feature_ts`	Spark Structured Streaming	~5 min	Real-time churn, recommendation
`transaction_features`	`user_id`	`feature_ts`	Scheduled batch (hourly)	~60 min	Churn, LTV prediction
`user_profile_features`	`user_id`	`updated_at`	CDC from OLTP (near real-time)	~2 min	All models
`product_features`	`product_id`	`feature_ts`	Scheduled batch (daily)	~24 hr	Recommendation, search ranking
`session_features`	`session_id`	`session_end_ts`	Streaming (micro-batch)	~1 min	Click-through rate, abandon prediction
`cohort_features`	`cohort_id`	`computed_at`	Weekly batch	~7 days	Segmentation, A/B analysis

Freshness vs cost tradeoff: Streaming features are ~10× more expensive to compute than batch features (continuous cluster vs scheduled job). Only promote a feature to streaming if your model's performance degrades meaningfully with stale data — validate this with an offline ablation study first.

Key Takeaways

Training-serving skew is the silent killer of production ML — the Feature Store eliminates it by encoding feature computation logic once and using it in both training and serving paths.
Point-in-time correct joins via timestamp_lookup_key are non-negotiable for any model trained on time-series data. A missing event_timestamp in your label table is a data leakage bug waiting to happen.
fe.log_model() is the right model logging call, not mlflow.sklearn.log_model(). It records feature lineage, enabling reproducible re-training and automatic feature lookup at serving time.
Watermarks in Structured Streaming are critical for stateful aggregations — without them, Spark accumulates state indefinitely and the job eventually OOMs. Set them to the maximum tolerable late-data window.
Online stores are only worth the operational cost when your SLA is under ~100ms. For batch scoring jobs or APIs with >500ms budgets, read directly from the offline Delta table.
fe.score_batch() is the cleanest way to run periodic batch inference — it handles PIT feature lookups automatically, keeps inference logic DRY, and logs results to Delta for downstream consumers.

References

Databricks — Feature Engineering in Unity Catalog (Overview)
🔗 https://docs.databricks.com/en/machine-learning/feature-store/uc/feature-tables-uc.html
Databricks — Create and Manage Online Tables
🔗 https://docs.databricks.com/en/machine-learning/feature-store/online-tables.html
Databricks — Point-in-Time Feature Lookups
🔗 https://docs.databricks.com/en/machine-learning/feature-store/time-series.html
Apache Spark — Structured Streaming Programming Guide
🔗 https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html
Apache Spark — Streaming Watermarks for Late Data Handling
🔗 https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html#handling-late-data-and-watermarking
Databricks — Feature Store Python API Reference
🔗 https://docs.databricks.com/en/machine-learning/feature-store/python-api.html
Databricks — Score Batch with Feature Store
🔗 https://docs.databricks.com/en/machine-learning/feature-store/score-batch.html
"Feature Stores for ML" — Feast Documentation (open-source reference)
🔗 https://docs.feast.dev/
"Rethinking Feature Stores" — Chip Huyen (huyenchip.com)
🔗 https://huyenchip.com/2023/01/08/feature-store.html
Databricks — Model Serving with Automatic Feature Lookup
🔗 https://docs.databricks.com/en/machine-learning/model-serving/feature-store-model-serving.html
"Building Machine Learning Pipelines" — Hannes Hapke & Catherine Nelson (O'Reilly)
🔗 https://www.oreilly.com/library/view/building-machine-learning/9781492053187/

This concludes the 4-part Databricks series:

Apache Spark Query Optimization on Databricks: Catalyst, AQE, and Photon Engine

Jubin Soni — Wed, 24 Jun 2026 09:32:37 +0000

A deep dive into how Spark transforms your SQL into a physical execution plan — and how Databricks layers Adaptive Query Execution and the Photon vectorized engine on top to squeeze out maximum performance.

Why Query Optimization Matters
The Catalyst Optimizer Pipeline
Stage 1: Parsing — From SQL to Unresolved Logical Plan
Stage 2: Analysis — Binding to the Catalog
Stage 3: Logical Optimization — Rule-Based Rewrites
Stage 4: Physical Planning — Strategies and Cost Models
Adaptive Query Execution (AQE)
The Photon Engine
Reading Explain Plans
Tuning Reference Table
References

Why Query Optimization Matters

A Spark query written by a human and a Spark query executed by the engine are often very different things. The gap between them — the optimization — is what separates a job that runs in 3 minutes from one that runs in 3 hours on identical hardware.

Databricks compounds Spark's native Catalyst optimizer with two additional layers:

Adaptive Query Execution (AQE) — re-optimizes the query at runtime using actual statistics collected mid-job
Photon — a C++ vectorized execution engine that replaces the JVM-based Spark executor for eligible operators

Understanding all three lets you write queries that cooperate with the engine rather than fight it.

The Catalyst Optimizer Pipeline

Catalyst is Spark's rule-based and cost-based query optimizer. Every query — whether written in SQL, DataFrame API, or Dataset API — passes through the same four-stage pipeline before a single byte of data is read.

Stage 1: Parsing — From SQL to Unresolved Logical Plan

# ── Catalyst Stage 1: Parsing ─────────────────────────────────────────────────
# Spark uses ANTLR4 to parse SQL into an Abstract Syntax Tree (AST).
# At this point column names are NOT validated — the plan is "unresolved".

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("catalyst-demo").getOrCreate()

# Both of these produce identical internal representations
df_api = (
    spark.table("prod.silver.events_clean")
        .filter("event_type = 'purchase'")
        .groupBy("platform")
        .agg({"revenue": "sum"})
)

sql_api = spark.sql("""
    SELECT platform, SUM(revenue) AS total_revenue
    FROM prod.silver.events_clean
    WHERE event_type = 'purchase'
    GROUP BY platform
""")

# Inspect the unresolved logical plan (before analysis)
df_api.explain(mode="formatted")
# Output includes:
# == Parsed Logical Plan ==
# 'Aggregate ['platform], ['platform, unresolvedAlias('sum('revenue), None)]
#   +- 'Filter ('event_type = 'purchase)
#      +- 'UnresolvedRelation [prod, silver, events_clean]

The key insight here: UnresolvedRelation and unresolvedAlias mean Spark hasn't touched the catalog yet. Column names could be typos at this point and Catalyst doesn't know.

Stage 2: Analysis — Binding to the Catalog

The Analyzer walks the unresolved AST and looks up every relation and attribute against the Catalog (in Databricks, this is Unity Catalog). It resolves column names, infers data types, validates references, and binds functions.

# ── Catalyst Stage 2: Analysis ────────────────────────────────────────────────
# After analysis, every column is resolved to a specific attribute with a type.
# AnalysisException is thrown HERE if a column doesn't exist.

from pyspark.sql import functions as F
from pyspark.sql.utils import AnalysisException

# Example of what Analysis catches:
try:
    spark.table("prod.silver.events_clean") \
        .select("nonexistent_column") \
        .show()
except AnalysisException as e:
    print(f"Analysis failed: {e}")
    # → AnalysisException: [UNRESOLVED_COLUMN.WITH_SUGGESTION]
    #   A column or function parameter with name `nonexistent_column` cannot be resolved.

# After successful analysis, inspect the resolved plan
df = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")
        .select("platform", "revenue", "user_id")
)

# The analyzed plan shows fully qualified attribute IDs like:
# == Analyzed Logical Plan ==
# platform: string, revenue: double, user_id: string
# Project [platform#42, revenue#67, user_id#31]
#   +- Filter (event_type#39 = purchase)
#      +- Relation prod.silver.events_clean[...] parquet
print(df._jdf.queryExecution().analyzed())

Stage 3: Logical Optimization — Rule-Based Rewrites

This is where Catalyst applies its ~100+ built-in rules to produce an equivalent but cheaper logical plan. Rules fire repeatedly in fixed-point iteration until the plan stabilises.

# ── Catalyst Stage 3: Key Optimization Rules ──────────────────────────────────

# RULE 1: Predicate Pushdown
# Catalyst moves filters as close to the data source as possible,
# so Spark reads fewer rows from Parquet.

df_before = (
    spark.table("prod.silver.events_clean")
        .join(
            spark.table("prod.silver.users_clean"),
            on="user_id"
        )
        .filter(F.col("event_type") == "purchase")  # ← filter AFTER join
)

# Catalyst rewrites this internally as if you wrote:
df_after_equivalent = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")  # ← filter BEFORE join
        .join(
            spark.table("prod.silver.users_clean"),
            on="user_id"
        )
)
# Result: potentially millions fewer rows shuffled during the join

# RULE 2: Column Pruning
# Catalyst removes columns not needed by downstream operators.
# Even if you select(*), Spark will only read the columns it needs.
df_pruned = (
    spark.table("prod.silver.events_clean")
        .select("*")
        .filter(F.col("event_type") == "purchase")
        .groupBy("platform")
        .agg(F.sum("revenue").alias("total_revenue"))
)
# Internally, Catalyst prunes all columns except: event_type, platform, revenue

# RULE 3: Constant Folding
# Expressions with only literals are evaluated at plan time, not per-row.
df_constants = spark.range(1000).select(
    F.lit(2 + 3 * 4).alias("always_14"),      # folded to Literal(14) at plan time
    F.col("id") * F.lit(1).alias("same_id"),   # simplified to just col("id")
)

# RULE 4: Boolean Simplification
# AND/OR chains with tautologies or contradictions are collapsed
df_simplified = spark.range(100).filter(
    (F.col("id") > 10) & F.lit(True)   # simplified to just (col("id") > 10)
)

# See all optimizations applied:
print(df_pruned._jdf.queryExecution().optimizedPlan())

Stage 4: Physical Planning — Strategies and Cost Models

The physical planner maps each logical operator to one or more physical implementations and selects the best one using a cost model. The most impactful decision here is join strategy selection.

# ── Catalyst Stage 4: Physical Planning & Join Strategies ────────────────────

# JOIN STRATEGY 1: Broadcast Hash Join (BHJ)
# Best when one side is small enough to fit in executor memory.
# No shuffle — the small table is broadcast to all workers.

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "10mb")  # default

large_df = spark.table("prod.silver.events_clean")      # 500GB
small_df = spark.table("prod.gold.product_catalog")     # 8MB ← will be broadcast

result_bhj = large_df.join(small_df, on="product_id")  # BHJ auto-selected

# Force BHJ with a broadcast hint (overrides threshold check):
from pyspark.sql.functions import broadcast
result_forced = large_df.join(broadcast(small_df), on="product_id")


# JOIN STRATEGY 2: Sort Merge Join (SMJ)
# Default for large-large joins. Both sides are sorted and merged.
# Requires a full shuffle — expensive but handles any size.

spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")  # disable BHJ
large_df2 = spark.table("prod.silver.transactions_clean")   # 200GB

result_smj = large_df.join(large_df2, on="user_id")  # SMJ selected


# JOIN STRATEGY 3: Shuffle Hash Join (SHJ)
# Hash-based, no sort. Chosen by AQE when one side is much smaller
# than the other but still above the broadcast threshold.
spark.conf.set("spark.sql.join.preferSortMergeJoin", "false")


# WHOLE-STAGE CODEGEN: Spark fuses multiple operators into a single
# Java function to avoid virtual dispatch overhead and intermediate objects.
# Verify it's active in your plan:
spark.conf.set("spark.sql.codegen.wholeStage", "true")  # default

result_bhj.explain(mode="formatted")
# Look for: *(1) BroadcastHashJoin — the *(N) prefix = WholeStageCodegen stage N

Adaptive Query Execution (AQE)

AQE is Databricks' most impactful runtime optimization layer. It materializes shuffle map output statistics at shuffle boundaries and uses them to make three key decisions after data has been partially processed.

# ── AQE Configuration ─────────────────────────────────────────────────────────

# AQE is ON by default in Databricks Runtime 7.3+
spark.conf.set("spark.sql.adaptive.enabled", "true")

# 1. Dynamic Partition Coalescing
# Merges small post-shuffle partitions to avoid thousands of tiny tasks
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled",       "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes",     "128mb")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionNum", "1")

# 2. Dynamic Join Strategy Switching
# Allows AQE to downgrade SMJ → BHJ at runtime if a side turns out small
spark.conf.set("spark.sql.adaptive.localShuffleReader.enabled",       "true")

# AQE broadcast threshold (can be higher than static threshold since
# we now KNOW the actual size)
spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold",       "30mb")

# 3. Skew Join Optimization
# Splits oversized partitions and replicates the non-skewed side
spark.conf.set("spark.sql.adaptive.skewJoin.enabled",                 "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor",   "5")    # 5x median
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256mb")

# Verify AQE decisions in the query plan:
df = (
    spark.table("prod.silver.events_clean")
        .join(spark.table("prod.silver.users_clean"), on="user_id")
        .groupBy("platform")
        .agg(F.sum("revenue").alias("total"))
)
df.explain(mode="formatted")
# Look for: AdaptiveSparkPlan isFinalPlan=true
# and: == Final Physical Plan == (shows post-AQE decisions)

The Photon Engine

Photon is Databricks' native vectorized query engine written in C++. It replaces the JVM-based Spark executor for eligible operations, processing data in column-oriented batches (vectors) rather than row-by-row.

# ── Photon Configuration & Verification ───────────────────────────────────────
# Photon is available on Databricks Runtime 9.1+ with Photon-enabled clusters.
# Enable it at the cluster level (UI: Cluster > Configuration > Enable Photon)
# or via config:

spark.conf.set("spark.databricks.photon.enabled", "true")

# Photon-accelerated operators (as of DBR 13.x):
# ✅ Scan (Parquet, Delta)         ✅ Filter / Project
# ✅ Hash Aggregate                ✅ Sort
# ✅ Broadcast Hash Join           ✅ Sort Merge Join
# ✅ Window functions              ✅ Union / Expand
# ✅ String functions              ✅ Math functions
# ❌ UDFs (Python/Scala)          ❌ Some complex types
# ❌ Streaming (partial)          ❌ RDD-based operations

# Verify Photon is executing your query:
df = spark.sql("""
    SELECT
        platform,
        DATE_TRUNC('month', event_ts)   AS month,
        SUM(revenue)                    AS total_revenue,
        COUNT(DISTINCT user_id)         AS unique_buyers,
        AVG(revenue)                    AS avg_order_value
    FROM prod.silver.events_clean
    WHERE event_type = 'purchase'
      AND event_ts >= '2024-01-01'
    GROUP BY platform, DATE_TRUNC('month', event_ts)
    ORDER BY month DESC, total_revenue DESC
""")

df.explain(mode="formatted")
# Look for operators prefixed with "Photon" in the physical plan:
# == Physical Plan ==
# PhotonResultStage
# +- PhotonSort [month DESC NULLS LAST, total_revenue DESC NULLS LAST]
#    +- PhotonShuffleExchangeSink hashpartitioning(platform, month)
#       +- PhotonGroupingAgg [platform, month], [sum(revenue), count(user_id), avg(revenue)]
#          +- PhotonFilter (event_type = purchase AND event_ts >= 2024-01-01)
#             +- PhotonScan parquet prod.silver.events_clean

# Photon performance metrics appear in Spark UI under "Photon Metrics":
# - Photon scan time
# - Photon total compute time
# - Rows processed by Photon vs fallback JVM

Reading Explain Plans

The explain(mode="formatted") output is your primary debugging tool. Here's how to read it efficiently:

# ── Explain Plan Modes ────────────────────────────────────────────────────────

df = (
    spark.table("prod.silver.events_clean")
        .filter(F.col("event_type") == "purchase")
        .join(broadcast(spark.table("prod.gold.product_catalog")), on="product_id")
        .groupBy("platform", "category")
        .agg(
            F.sum("revenue").alias("total_revenue"),
            F.count("*").alias("transaction_count")
        )
)

# Mode 1: simple (default) — compact tree
df.explain()

# Mode 2: extended — all 4 plan stages side by side
df.explain(mode="extended")

# Mode 3: formatted — human-readable with operator details (RECOMMENDED)
df.explain(mode="formatted")

# Mode 4: cost — includes estimated row counts and sizes (requires ANALYZE TABLE)
df.explain(mode="cost")

# Mode 5: codegen — shows generated Java code for WholeStageCodegen
df.explain(mode="codegen")


# ── Key Signals to Look For ───────────────────────────────────────────────────

# ✅ GOOD signs:
# *(N) prefix          → WholeStageCodegen active (operators fused)
# BroadcastHashJoin    → small table correctly broadcast, no shuffle
# PhotonXxx            → Photon accelerating this operator
# AdaptiveSparkPlan    → AQE is engaged
# PartitionFilters     → Delta/Parquet file skipping active
# PushedFilters        → filters pushed to Parquet reader

# ❌ WARNING signs:
# Exchange (shuffle)   → unexpected shuffle (missing broadcast hint?)
# SortMergeJoin        → large-large join (may need Z-ORDER or AQE tuning)
# HashAggregate x2     → partial + final agg = shuffle involved
# CartesianProduct     → missing join condition! Will OOM on large tables
# ObjectHashAggregate  → non-codegen path, JVM overhead
# GenerateXxx          → explode() or similar, can't be fused


# ── ANALYZE TABLE: feed statistics to CBO ─────────────────────────────────────
# Without stats, Catalyst uses default estimates (1M rows, 8 bytes/col).
# Run ANALYZE to give the Cost-Based Optimizer real numbers.

spark.sql("ANALYZE TABLE prod.silver.events_clean COMPUTE STATISTICS")
spark.sql("""
    ANALYZE TABLE prod.silver.events_clean
    COMPUTE STATISTICS FOR COLUMNS user_id, event_type, platform, revenue
""")
# Now explain(mode="cost") shows real row counts and sizes

Tuning Reference Table

A quick-reference guide for the most impactful Spark/Databricks configs, what they control, and when to change them:

Config Key	Default	What It Controls	When to Tune
`spark.sql.adaptive.enabled`	`true`	Master AQE switch	Keep on; only disable for debugging
`spark.sql.adaptive.advisoryPartitionSizeInBytes`	`64mb`	Target post-coalesce partition size	Increase to `128mb`–`256mb` for large shuffles
`spark.sql.adaptive.skewJoin.enabled`	`true`	AQE skew split	Keep on; tune `skewedPartitionFactor` if needed
`spark.sql.autoBroadcastJoinThreshold`	`10mb`	Static BHJ threshold	Increase to `50mb`–`100mb` if executor memory allows
`spark.sql.adaptive.autoBroadcastJoinThreshold`	`30mb`	AQE runtime BHJ threshold	Increase if AQE isn't catching small tables
`spark.sql.shuffle.partitions`	`200`	Default shuffle partition count	Set to `8 × num_cores` for your cluster
`spark.sql.files.maxPartitionBytes`	`128mb`	Max bytes per Parquet read partition	Reduce for high-parallelism scans
`spark.databricks.photon.enabled`	`true`	Photon vectorized engine	Keep on; disable only for UDF-heavy jobs
`spark.sql.codegen.wholeStage`	`true`	Whole-Stage CodeGen fusion	Keep on; disable only for debugging
`spark.sql.statistics.histogram.enabled`	`false`	Column histograms for CBO	Enable after running ANALYZE TABLE
`spark.sql.cbo.enabled`	`true`	Cost-Based Optimizer	Keep on; requires ANALYZE TABLE to be useful
`spark.databricks.delta.optimizeWrite.enabled`	`true`	Auto bin-pack write files	Keep on for all Delta writes

Key Takeaways

Catalyst has four stages: Parse → Analyze → Optimize → Plan. Each stage has a distinct job, and understanding them tells you exactly where to look when a query misbehaves.
Predicate pushdown and column pruning are the two most impactful automatic optimizations — they reduce the data volume Spark has to move before any aggregation or join.
AQE is not a set-and-forget feature: tune advisoryPartitionSizeInBytes to your actual data sizes, and verify its decisions with explain(mode="formatted") — look for AdaptiveSparkPlan isFinalPlan=true.
Photon drops in transparently for most SQL and DataFrame operations. The exceptions are Python UDFs, RDD operations, and some complex types — refactor these away from hot paths.
Run ANALYZE TABLE ... COMPUTE STATISTICS FOR COLUMNS on your most-joined tables. The CBO's join ordering and strategy decisions improve dramatically with real statistics vs. default estimates.
explain(mode="formatted") is your most important debugging tool — learn to read it before reaching for cluster config changes.

References

Apache Spark — Catalyst Optimizer (Deep Dive Paper, Armbrust et al., SIGMOD 2015)
🔗 https://people.csail.mit.edu/matei/papers/2015/sigmod_spark_sql.pdf
Databricks — Adaptive Query Execution
🔗 https://docs.databricks.com/en/optimizations/aqe.html
Apache Spark Docs — Adaptive Query Execution
🔗 https://spark.apache.org/docs/latest/sql-performance-tuning.html#adaptive-query-execution
Databricks — Photon Runtime
🔗 https://docs.databricks.com/en/compute/photon.html
Databricks Blog — Photon: A Fast Query Engine for Lakehouse Systems
🔗 https://www.databricks.com/blog/2021/06/17/photon-a-fast-query-engine-for-lakehouse-systems.html
Databricks — Cost-Based Optimizer
🔗 https://docs.databricks.com/en/optimizations/cost-based-optimizer.html
Apache Spark — Performance Tuning Guide
🔗 https://spark.apache.org/docs/latest/sql-performance-tuning.html
Databricks — Broadcast Join Hints
🔗 https://docs.databricks.com/en/optimizations/join-reorder.html
"Photon: A Fast Query Engine for Lakehouse Systems" (Behm et al., SIGMOD 2022)
🔗 https://dl.acm.org/doi/10.1145/3514221.3526054
Spark by Examples — Explain Plan Modes
🔗 https://sparkbyexamples.com/spark/spark-explain-plan-modes/

Next in the series:

Previous in the series