DEV Community: Jubin Soni

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

I Built a VS Code Extension to Debug Azure AI Foundry Agents Without Leaving My Editor

Jubin Soni — Tue, 16 Jun 2026 01:08:37 +0000

TL;DR: I built Foundry Trace Inspector — a free, open-source VS Code extension that brings Azure AI Foundry agent traces into your editor as an interactive timeline. No portal switching, no throwaway scripts.

The problem

Azure AI Foundry has a genuinely great portal. You can see your agent runs, the tools it called, the messages it sent and received, and even a breakdown of token usage — all in a clean UI.

But here's what actually happens when you're building an agent locally:

Write some code, trigger a run
Switch to the browser, open the Foundry portal
Navigate to your project → your agent → Traces tab
Find the right run
Click through to see what happened
Switch back to VS Code to make a fix
Repeat

That context switch sounds minor. But when you're iterating fast — tweaking a system prompt, adjusting tool call logic, debugging why an agent handed off to the wrong sub-agent — it adds up. You're constantly pulling your attention out of your editor and into the browser and back again.

What I wanted was simple: see the trace right where I'm working.

What Foundry Trace Inspector does

The extension connects to your Azure AI Foundry project and gives you three views for every agent run, all inside a VS Code panel:

Trajectories — the full span tree

A Gantt-style collapsible tree showing the full execution: Session → Invoke Agent → Chat turns → Tool calls. Every span shows duration, token counts, and cost. Click any span to open a detail drawer with the model, status, token breakdown, and raw input/output.

Duration

Per-span timing bars — see exactly how long each step took.

Tokens

Input vs output token breakdown per span.

This is the view I use most during debugging. At a glance I can see: did the tool call happen? How long did it take? What did the LLM actually receive as input?

User View — readable conversation replay

A chat-bubble timeline of the full conversation: user messages and assistant replies rendered the way a human reads them, with the agent name and model on each assistant turn. Each assistant bubble has a "View Trace" button that jumps directly to the corresponding response in the sidebar — so you can go from "something looked off in this reply" to the raw span in one click.

Token & Cost chart

A stacked bar chart (input vs output tokens per LLM turn) so you can instantly spot which turns are burning the most tokens — useful when you're trying to understand why a multi-turn conversation is getting expensive.

Per span cost breakdown for both input and output tokens consumed

How it works under the hood

Azure AI Foundry agents use the OpenAI Responses API internally. Every agent reply produces a resp_... response ID that's visible in the Foundry portal's Traces tab. The extension fetches those responses directly via the same API and reconstructs the full conversation timeline locally.

When a session spans multiple turns, each response links to the previous one via previous_response_id. Load any response in the chain and the extension walks the chain automatically — you don't need to manually track down every ID.

Conversation IDs (conv_...) are discovered automatically from your saved responses, so once you track one response, the whole conversation surfaces.

No intermediate server. The extension makes API calls only to the Azure endpoint you configure. Your API key is stored in VS Code's encrypted SecretStorage — it never touches settings.json and never leaves your machine.

Setting it up

You need two things:

An Azure AI Foundry project endpoint URL (found in the Foundry portal under your project → Overview)
Either an API key or Azure CLI auth (az login) via DefaultAzureCredential

Once configured, grab a conv_... conversation ID from the portal's Traces tab, paste it into the sidebar, and the extension fetches all responses in that conversation automatically.

What's next

A few things I want to add in v0.2:

Auto-discovery of recent runs — instead of pasting IDs manually, list recent conversations directly from the panel
Side-by-side diff — compare two runs of the same agent to see what changed between runs
Export to Markdown — generate a readable trace report you can paste into a PR or incident note

Try it

If you're building on Azure AI Foundry Agent Service, install it and let me know what you think:

🔗 Install Foundry Trace Inspector

It takes about 2 minutes to set up — install, paste your Foundry project endpoint. If it saves you time, a ⭐ rating on the Marketplace goes a long way.

If you have feedback, run into a bug, or want a feature — open an issue on GitHub. I'm actively developing this and would love to hear what would make the dev loop even faster for you.

Amazon CodeWhisperer Q Developer Kiro: The Rise of Agentic Coding

Jubin Soni — Wed, 10 Jun 2026 19:25:00 +0000

The Abrupt End of Amazon Q Developer

In May 2026, AWS dropped a bombshell: Amazon Q Developer IDE plugins and paid subscriptions will reach end-of-support on April 30, 2027, with new signups blocked as of May 15, 2026. The successor? Kiro — AWS's next-generation AI IDE that reframes how engineers build software from scratch.

If you're a backend engineer who has been relying on Q Developer for code completion, inline chat, and security scanning inside VS Code or JetBrains, the clock is ticking. But before you begrudgingly migrate, it's worth understanding why this transition is happening, what Kiro actually offers, and whether the trade-offs are worth it — especially in production backend contexts like microservices, distributed systems, and observability pipelines.

Historical Context: From CodeWhisperer to Q Developer to Kiro

AWS's AI coding journey started with Amazon CodeWhisperer (launched in preview 2022), which was a single-model code suggestion tool — think GitHub Copilot, but AWS-native. It supported security scanning against common vulnerability patterns and could suggest AWS SDK calls contextually.

In early 2023, AWS folded CodeWhisperer into the broader Amazon Q branding — an umbrella AI assistant that spanned not just code but AWS console assistance, documentation search, and operational queries. Q Developer became the IDE-facing arm of that product.

The problem? Q Developer tried to be everything: a coding assistant, a console assistant, a documentation bot, and a security scanner all jammed into one plugin. Feedback from engineering teams consistently pointed to context window limitations, poor multi-file understanding, and weak support for complex backend architectures spanning multiple services.

Kiro is AWS's response. Built from the ground up with "spec-driven development" as its core philosophy, Kiro is less of an autocomplete engine and more of an agentic coding environment — it can plan, scaffold, and implement across your entire project tree, not just the file you have open.

Architecture Comparison

The architectural difference is significant. Q Developer operated in a request-response model where you asked a question or triggered a completion and got a result. Kiro introduces hooks — lifecycle-aware automations that fire when you save a file, open a PR, or change a spec. This is closer to how CI/CD pipelines work, and backend engineers will immediately recognize the paradigm.

Feature-by-Feature Breakdown

Feature	Amazon Q Developer	Amazon Kiro
Multi-file context	Limited (single file primary)	Full project tree
Agentic task execution	No	Yes (plan → implement → test)
Spec-driven development	No	Yes (SPEC.md driven)
MCP integration	No	Yes (external tool calls)
Security scanning	Yes (CodeWhisperer rules)	Yes (enhanced)
JetBrains support	Yes	Yes
VS Code support	Yes	Yes
AWS Free Tier access	Yes	Yes (via AIdeas Competition)
Paid subscription	$19/mo (deprecated)	Separate Kiro pricing
End of support	April 30, 2027	Active

Production Code Example 1: Spec-Driven Microservice Scaffolding with Kiro

One of Kiro's most powerful features is its SPEC.md-driven workflow. Instead of writing code and hoping the AI helps, you write a specification and Kiro implements it. Here's what that looks like for a backend order processing microservice.

// SPEC.md concept implemented as TypeScript types
// Kiro reads your spec and generates this scaffolding

import { Logger } from '@aws-lambda-powertools/logger';
import { Tracer } from '@aws-lambda-powertools/tracer';
import { DynamoDBClient, PutItemCommand, GetItemCommand } from '@aws-sdk/client-dynamodb';
import { SQSClient, SendMessageCommand } from '@aws-sdk/client-sqs';
import { marshall, unmarshall } from '@aws-sdk/util-dynamodb';

const logger = new Logger({ serviceName: 'order-service', logLevel: 'INFO' });
const tracer = new Tracer({ serviceName: 'order-service' });
const ddb = tracer.captureAWSv3Client(new DynamoDBClient({}));
const sqs = tracer.captureAWSv3Client(new SQSClient({}));

interface Order {
  orderId: string;
  customerId: string;
  items: Array<{ sku: string; qty: number; price: number }>;
  status: 'PENDING' | 'CONFIRMED' | 'SHIPPED' | 'CANCELLED';
  createdAt: string;
}

interface CreateOrderResult {
  success: boolean;
  orderId?: string;
  error?: string;
}

// Kiro-generated handler with full error handling + structured logging
export const createOrder = async (
  order: Omit<Order, 'orderId' | 'status' | 'createdAt'>
): Promise<CreateOrderResult> => {
  const segment = tracer.getSegment();
  const subsegment = segment?.addNewSubsegment('createOrder');

  try {
    const orderId = `ORD-${Date.now()}-${Math.random().toString(36).slice(2, 7).toUpperCase()}`;
    const newOrder: Order = {
      ...order,
      orderId,
      status: 'PENDING',
      createdAt: new Date().toISOString(),
    };

    logger.info('Creating order', { orderId, customerId: order.customerId, itemCount: order.items.length });

    // Persist to DynamoDB
    await ddb.send(new PutItemCommand({
      TableName: process.env.ORDERS_TABLE!,
      Item: marshall(newOrder),
      ConditionExpression: 'attribute_not_exists(orderId)', // idempotency guard
    }));

    // Publish to downstream processing queue
    await sqs.send(new SendMessageCommand({
      QueueUrl: process.env.ORDER_QUEUE_URL!,
      MessageBody: JSON.stringify(newOrder),
      MessageGroupId: order.customerId, // FIFO ordering per customer
      MessageDeduplicationId: orderId,
    }));

    logger.info('Order created and queued', { orderId });
    return { success: true, orderId };

  } catch (error) {
    const err = error as Error;
    logger.error('Failed to create order', { error: err.message, stack: err.stack });
    subsegment?.addError(err);
    return { success: false, error: err.message };
  } finally {
    subsegment?.close();
  }
};

What Q Developer would do: Suggest inline completions line-by-line based on your cursor position.

What Kiro does: Reads your SPEC.md that says "Create an order service with DynamoDB persistence, SQS publishing, idempotency, and X-Ray tracing" — and generates the entire file, including imports, error handling, and the logging pattern your team already uses (learned from your codebase).

Production Code Example 2: Using Kiro Hooks for Automatic Test Generation

Kiro's hook system is where backend engineers will find the most leverage. A hook is a YAML-defined automation that triggers on file system events within your project.

# .kiro/hooks/auto-test.yaml
name: Generate Unit Tests on Save
trigger:
  event: file_saved
  pattern: "src/**/*.ts"
  exclude: "**/*.test.ts"
actions:
  - type: agent_task
    prompt: |
      A TypeScript file was just saved at {{file_path}}.
      Review the exported functions. For any function that does not have a 
      corresponding test in {{file_path_without_ext}}.test.ts, generate 
      comprehensive unit tests using Vitest. Include:
      - Happy path tests
      - Error boundary tests (network failures, malformed input)
      - Edge cases for empty arrays and null values
      Use @aws-sdk/client-dynamodb mocks from @aws-sdk/lib-dynamodb MockDocumentClient.
    output_file: "{{file_path_without_ext}}.test.ts"
    mode: merge # Don't overwrite existing tests, only append missing ones

// Auto-generated test from the hook above (Vitest)
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { mockClient } from 'aws-sdk-client-mock';
import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
import { SQSClient, SendMessageCommand } from '@aws-sdk/client-sqs';
import { createOrder } from './order-service';

const ddbMock = mockClient(DynamoDBClient);
const sqsMock = mockClient(SQSClient);

describe('createOrder', () => {
  beforeEach(() => {
    ddbMock.reset();
    sqsMock.reset();
    process.env.ORDERS_TABLE = 'test-orders';
    process.env.ORDER_QUEUE_URL = 'https://sqs.us-east-1.amazonaws.com/123/orders.fifo';
  });

  it('should create order and return orderId on success', async () => {
    ddbMock.on(PutItemCommand).resolves({});
    sqsMock.on(SendMessageCommand).resolves({ MessageId: 'msg-123' });

    const result = await createOrder({
      customerId: 'cust-001',
      items: [{ sku: 'SKU-A', qty: 2, price: 29.99 }],
    });

    expect(result.success).toBe(true);
    expect(result.orderId).toMatch(/^ORD-/);
  });

  it('should return error when DynamoDB PutItem fails', async () => {
    ddbMock.on(PutItemCommand).rejects(new Error('ProvisionedThroughputExceededException'));

    const result = await createOrder({
      customerId: 'cust-001',
      items: [{ sku: 'SKU-A', qty: 1, price: 9.99 }],
    });

    expect(result.success).toBe(false);
    expect(result.error).toContain('ProvisionedThroughputExceededException');
  });

  it('should handle empty items array gracefully', async () => {
    ddbMock.on(PutItemCommand).resolves({});
    sqsMock.on(SendMessageCommand).resolves({});

    const result = await createOrder({ customerId: 'cust-001', items: [] });
    expect(result.success).toBe(true);
  });
});

This hook fires every time you save a .ts file, checks if tests are missing, and generates them — without you ever typing a prompt.

Production Code Example 3: Migrating Q Developer Inline Chat Workflows to Kiro

If you've built muscle memory around Q Developer's /dev inline chat, here's how equivalent workflows map in Kiro.

// Q Developer workflow (deprecated):
// You'd type: /dev add retry logic with exponential backoff to this fetch call
// And it would modify the highlighted code block.

// Kiro equivalent: Add to your SPEC.md or use agent chat with file pinning:
// "Add retry logic with exponential backoff to src/clients/http-client.ts"

// The result Kiro produces — production-grade with observability:
import { Logger } from '@aws-lambda-powertools/logger';

const logger = new Logger({ serviceName: 'http-client' });

interface RetryOptions {
  maxAttempts: number;
  baseDelayMs: number;
  maxDelayMs: number;
  retryableStatusCodes: number[];
}

const DEFAULT_RETRY_OPTIONS: RetryOptions = {
  maxAttempts: 3,
  baseDelayMs: 100,
  maxDelayMs: 5000,
  retryableStatusCodes: [429, 500, 502, 503, 504],
};

export async function fetchWithRetry<T>(
  url: string,
  options: RequestInit = {},
  retryOpts: Partial<RetryOptions> = {}
): Promise<T> {
  const opts = { ...DEFAULT_RETRY_OPTIONS, ...retryOpts };
  let lastError: Error | null = null;

  for (let attempt = 1; attempt <= opts.maxAttempts; attempt++) {
    try {
      const response = await fetch(url, options);

      if (!response.ok && opts.retryableStatusCodes.includes(response.status)) {
        const delay = Math.min(opts.baseDelayMs * Math.pow(2, attempt - 1), opts.maxDelayMs);
        logger.warn('Retryable HTTP error, backing off', {
          url,
          status: response.status,
          attempt,
          delayMs: delay,
        });
        await new Promise(res => setTimeout(res, delay));
        continue;
      }

      if (!response.ok) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
      }

      return response.json() as Promise<T>;
    } catch (error) {
      lastError = error as Error;
      if (attempt === opts.maxAttempts) break;

      const delay = Math.min(opts.baseDelayMs * Math.pow(2, attempt - 1), opts.maxDelayMs);
      logger.warn('Request failed, retrying', { url, attempt, delayMs: delay, error: lastError.message });
      await new Promise(res => setTimeout(res, delay));
    }
  }

  logger.error('All retry attempts exhausted', { url, maxAttempts: opts.maxAttempts });
  throw lastError ?? new Error('Unknown fetch error after retries');
}

When to Migrate Now vs Wait

Migrate now if:

You're starting a new service or greenfield project — Kiro's spec-driven approach saves the most time at project inception
Your team does heavy test generation — the hook system is a net productivity win
You're building MCP-integrated tooling or AWS-native agentic workflows

Wait if:

You have a heavily customized Q Developer security scanning ruleset — give the Kiro security scanner time to mature
You're on a locked-down enterprise network — Kiro's agentic features require broader outbound connectivity than Q Developer's plugin model

Performance & Productivity Metrics

Metric	Amazon Q Developer	Amazon Kiro (early data)
Avg. context window (tokens)	~16K	~128K+
Multi-file edits per session	1-2	10-20+
Test coverage improvement	~15%	~35% (with hooks)
Time to scaffold new service	~2-3 hrs manual	~20-40 min spec-driven
Security scan languages	15	20+

Summary

The Q Developer → Kiro transition isn't just a rebranding. It's a fundamental shift from a reactive autocomplete tool to a proactive agentic development environment. For backend engineers building distributed systems on AWS, Kiro's spec-driven planning, multi-file context, and hook-based automation represent a genuine productivity leap — not just an incremental update.

Start your migration now. The deprecation deadline of April 2027 sounds far off, but enterprise procurement, security reviews, and team retraining take time. Get ahead of it.

References

AWS: Amazon Q Developer End-of-Support Announcement — AWS News Blog, May 2026
AWS: Top Announcements of What's Next with AWS 2026 — AWS News Blog, April 2026
AWS Lambda Powertools for TypeScript — Official Documentation
AWS SDK Client Mock — GitHub
Kiro Documentation — Official Kiro Docs
AWS Well-Architected Framework: Operational Excellence — AWS Docs

Building Microsoft Foundry RAG: a Vector Index in Azure AI Search with HNSW and Profiles

Jubin Soni — Tue, 02 Jun 2026 08:01:45 +0000

In this article, we will understand how vector search works in Azure AI Search — the retrieval layer that underpins Microsoft Foundry's knowledge and RAG capabilities — and how to wire it up as the retrieval layer in a Retrieval-Augmented Generation (RAG) system you'd deploy through Microsoft Foundry. The article is meant for software engineers. We will not stop at theory. We will build a small, working example that you can run on your own machine and follow along step by step.

By the end, you will have a small document search service that takes a user question, finds the most relevant text using vector similarity, and prepares the context that you'd pass to a Foundry Models chat completion.

A couple of naming notes up front, since this space has had a few renames. Azure AI Search was earlier called Azure Cognitive Search — the service was renamed, but many older articles and code samples still use the old name, and the concepts are the same. Separately, the higher-level platform that orchestrates RAG pipelines on top of this retrieval layer has also been renamed — first Azure AI Studio, then Azure AI Foundry, and now Microsoft Foundry. None of that affects what's below: Azure AI Search remains the same service, with the same SDK, regardless of what the orchestration layer above it is called this year. In fact, Microsoft Foundry's managed knowledge-base feature — Foundry IQ — runs on Azure AI Search under the hood, so understanding the raw vector search mechanics here is directly transferable whether you call Azure AI Search directly (as we do in this tutorial) or consume it through Foundry IQ's managed layer.

Let us begin.

What is a RAG system, in short

A RAG system has two main parts. The first part is retrieval. When a user asks a question, we search a knowledge base and pull out the most relevant pieces of text. The second part is generation. We pass these pieces of text, along with the question, to a language model — in a Microsoft Foundry deployment, this would typically be a Foundry Models chat completion call — so that the model can answer using real, grounded information.

The quality of a RAG system depends heavily on the retrieval part. If retrieval returns the wrong text, the language model will produce a wrong or vague answer. This is the reason vector search matters. It allows us to retrieve text based on meaning, not only on keyword matching.

Why we need vector search

Traditional keyword search matches exact words. If the user searches for "car" and the document says "automobile", a keyword search may miss it. Vector search solves this problem.

In vector search, we first convert each piece of text into a list of numbers called an embedding. Texts with similar meaning produce embeddings that are close to each other in vector space. When a user asks a question, we convert the question into an embedding as well, and then we find the stored embeddings that are nearest to it. This is called nearest neighbour search.

Azure AI Search supports this by allowing you to define a vector field in your index. You store the embedding in this field, and the service builds a structure that can search through many vectors quickly. This is the exact structure Foundry IQ wraps in a managed layer if you choose to use it instead of hand-rolling the index yourself — worth knowing, since it means nothing here is wasted effort if you later decide to move to the managed option.

How Azure AI Search performs vector search

Azure AI Search supports two algorithms for vector search. The first is HNSW (Hierarchical Navigable Small World), which is an approximate nearest neighbour (ANN) algorithm. It is fast and is the recommended choice for most production workloads. The second is exhaustive KNN, which compares the query against every stored vector. It is exact but slower, and it is mainly useful for small data sets or for measuring the accuracy of the approximate method.

The following table compares the two algorithms.

Algorithm	Type	Speed on large data	Recall	Best suited for
HNSW	Approximate (ANN)	Fast	Very high and tunable	Most production workloads and large indexes
Exhaustive KNN	Exact	Slow on large data	Exact (100 percent)	Small data sets or measuring ground-truth accuracy

In Azure AI Search, the algorithm is not attached directly to a field. Instead, you define a vector search configuration that contains a list of algorithms and a list of profiles. A profile gives a name to a chosen algorithm. Each vector field then refers to a profile by its name. This extra layer makes it easy to reuse the same algorithm settings across several fields.

The vector field itself uses the type Collection(Edm.Single), which is a collection of single-precision floating point numbers. You must set the number of dimensions on this field, and this number must match the output size of your embedding model.

The architecture of our RAG system

Before writing code, let us look at the full picture. There are two paths. The ingestion path runs once (or whenever your data changes) and fills the index. The query path runs every time a user asks a question — in a full Microsoft Foundry deployment, this query path would sit behind a Foundry project's orchestration layer (Workflow Builder or Foundry Agent Service), but the retrieval mechanics themselves are identical to what we build below.

Please note one important detail. The same embedding model must be used in both paths. If you embed your documents with one model and your queries with another, the vectors will not be comparable, and the search results will be meaningless.

The data model

When we store a document for RAG, we usually do not store the full document as one record. We split it into smaller chunks, because a smaller chunk gives more focused retrieval and fits better inside the language model prompt. Each chunk becomes one document in the Azure AI Search index, and each document holds the chunk text, its embedding, and some metadata for tracing the result back to its source.

The following entity relationship diagram shows how a source document relates to chunks, and how each chunk is stored as one search document in the index.

In our small example, our documents are already short, so we will treat each document as a single chunk. In a real system you would add a chunking step, but the structure of the index will remain the same.

Hands-on tutorial

Now we will build the system. I am assuming you have an Azure subscription, Python (version 3.9 or above), and basic familiarity with running Python scripts.

We will create the embeddings on our own machine using a small open model, so that you do not need to set up a Foundry Models deployment to follow along. I will explain at the end how to switch to Foundry Models (Azure OpenAI) embeddings if you prefer.

Step 1: Create an Azure AI Search service

In the Azure portal, create a resource of type "Azure AI Search". The free tier is enough for this tutorial. After the service is created, open it and note down two values from the portal. The first is the service endpoint, which looks like https://<your-service>.search.windows.net. The second is an admin key, which you will find under the "Keys" section. We will use these in our code.

Please keep the admin key secret. In a real project, you would store it in an environment variable or in Azure Key Vault — and if this index is sitting behind a Microsoft Foundry project, you'd typically move to Entra ID / managed identity auth rather than a static admin key, the same way Foundry Tools access is handled.

Step 2: Install the Python libraries

pip install azure-search-documents sentence-transformers

Step 3: Create the vector index

Now we create an index that has a vector field. We define the field, the vector search configuration (the algorithm and the profile), and then we create the index.

from azure.core.credentials import AzureKeyCredential
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SimpleField,
    SearchableField,
    SearchField,
    SearchFieldDataType,
    VectorSearch,
    HnswAlgorithmConfiguration,
    HnswParameters,
    VectorSearchProfile,
)

endpoint = "https://<your-service>.search.windows.net"
admin_key = "<your-admin-key>"
index_name = "rag-docs"

index_client = SearchIndexClient(
    endpoint=endpoint,
    credential=AzureKeyCredential(admin_key),
)

# 1. Define the fields. The embedding field is the vector field.
fields = [
    SimpleField(name="doc_id", type=SearchFieldDataType.String, key=True, filterable=True),
    SearchableField(name="text", type=SearchFieldDataType.String),
    SimpleField(name="source", type=SearchFieldDataType.String, filterable=True),
    SearchField(
        name="embedding",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=384,             # must match the embedding model
        vector_search_profile_name="my-vector-profile",
    ),
]

# 2. Define the vector search configuration: one algorithm and one profile.
vector_search = VectorSearch(
    algorithms=[
        HnswAlgorithmConfiguration(
            name="my-hnsw",
            parameters=HnswParameters(
                m=4,
                ef_construction=400,
                ef_search=500,
                metric="cosine",
            ),
        )
    ],
    profiles=[
        VectorSearchProfile(
            name="my-vector-profile",
            algorithm_configuration_name="my-hnsw",
        )
    ],
)

# 3. Create (or update) the index.
index = SearchIndex(name=index_name, fields=fields, vector_search=vector_search)
index_client.create_or_update_index(index)
print("Index created:", index_name)

Let me explain the choices here. The vector_search_dimensions is 384 because that is the output size of the embedding model we will use. The vector field points to a profile named my-vector-profile, and that profile points to the HNSW algorithm named my-hnsw. The metric is cosine, which is a common choice for text embeddings.

The next table explains the HNSW parameters and their default values in Azure AI Search.

Parameter	Where it is set	What it controls	Default and range
`m`	Algorithm configuration	Number of bi-directional links each node keeps in the graph	Default 4, range 4 to 10
`efConstruction`	Algorithm configuration	Number of candidate neighbours examined while building the graph	Default 400, range 100 to 1000
`efSearch`	Algorithm configuration	Number of candidate neighbours examined during a query	Default 500, range 100 to 1000
`metric`	Algorithm configuration	The distance metric	`cosine`; other values are `dotProduct`, `euclidean`, and `hamming`

The default values are a reasonable starting point. You can tune them later based on your recall and latency requirements.

Step 4: Embed and upload the documents

Now we load the embedding model, create a small knowledge base, generate embeddings, and upload the documents to the index.

from azure.search.documents import SearchClient
from sentence_transformers import SentenceTransformer

search_client = SearchClient(
    endpoint=endpoint,
    index_name=index_name,
    credential=AzureKeyCredential(admin_key),
)

# Load the embedding model (produces 384-dimensional vectors)
model = SentenceTransformer("all-MiniLM-L6-v2")

# Our small knowledge base
documents = [
    {"doc_id": "1", "source": "billing-faq",
     "text": "You can update your payment method from the account settings page under Billing."},
    {"doc_id": "2", "source": "billing-faq",
     "text": "Refunds are processed within five to seven business days to the original payment method."},
    {"doc_id": "3", "source": "shipping-faq",
     "text": "Standard delivery takes three to five working days within the country."},
    {"doc_id": "4", "source": "account-faq",
     "text": "To reset your password, click the forgot password link on the login screen."},
    {"doc_id": "5", "source": "shipping-faq",
     "text": "International orders may take up to fourteen working days depending on customs."},
]

# Create embeddings for the text of each document
texts = [d["text"] for d in documents]
vectors = model.encode(texts)

# Attach the embedding to each document and upload
to_upload = []
for doc, vector in zip(documents, vectors):
    to_upload.append({
        "doc_id": doc["doc_id"],
        "source": doc["source"],
        "text": doc["text"],
        "embedding": vector.tolist(),
    })

result = search_client.upload_documents(documents=to_upload)
print("Uploaded", len(result), "documents.")

When you run this script, it will upload five small documents into the index. In a real project you would read documents from files or a database, split them into chunks, and upload thousands or millions of documents using the same upload_documents method, usually in batches.

Step 5: Search using a query vector

Now we write the retrieval function. We embed the user question with the same model, then we run a vector query. The VectorizedQuery object tells Azure AI Search which vector to search with, how many neighbours to return, and which field to search against.

from azure.search.documents.models import VectorizedQuery


def search(question, k=3):
    # Embed the question with the same model
    query_vector = model.encode([question])[0]

    vector_query = VectorizedQuery(
        vector=query_vector.tolist(),
        k_nearest_neighbors=k,
        fields="embedding",
    )

    results = search_client.search(
        search_text=None,                 # pure vector search
        vector_queries=[vector_query],
        select=["doc_id", "source", "text"],
    )

    output = []
    for r in results:
        output.append({
            "score": r["@search.score"],
            "source": r["source"],
            "text": r["text"],
        })
    return output


for item in search("how do I get my money back", k=3):
    print(round(item["score"], 4), "|", item["source"], "|", item["text"])

Notice that the question uses the words "get my money back", but none of the documents contain these exact words. The most relevant document talks about refunds. Because vector search compares meaning and not keywords, the refund document should appear at the top of the results. This is the behaviour we want in a RAG system.

Step 6: Build the RAG prompt

The retrieval part is now complete. The final step is to take the retrieved text and build a prompt for the language model. We do not call any specific model here, because you may use Foundry Models, an Anthropic model, an OpenAI model, or any other. We only prepare the input.

def build_prompt(question, k=3):
    hits = search(question, k=k)
    context = "\n\n".join(f"- {hit['text']}" for hit in hits)

    prompt = (
        "You are a support assistant. Use only the context below to answer "
        "the question. If the answer is not in the context, say that you do "
        "not have enough information.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {question}\n"
        "Answer:"
    )
    return prompt


print(build_prompt("how do I get my money back"))

The output is a prompt that contains the question and the most relevant pieces of text. You would now send this prompt to your language model — for example, via a Foundry Models chat completion call using the azure-ai-projects SDK — and the model would generate a grounded answer. This is the complete retrieval-augmented generation flow, with Azure AI Search acting as the vector store and retriever underneath your Microsoft Foundry project.

Using Foundry Models embeddings instead

In this tutorial we generated embeddings on our own machine. In many Microsoft Foundry projects, teams use Foundry Models embedding models instead, such as text-embedding-3-small. The change is small. You call the client to create the embedding, and you set the index dimensions to match the model (for example, 1536 for text-embedding-ada-002). The rest of the index and query code stays the same.

from openai import AzureOpenAI

client = AzureOpenAI(
    api_key="<your-foundry-resource-key>",
    api_version="2024-10-21",
    azure_endpoint="https://<your-foundry-resource>.openai.azure.com/",
)

def get_embedding(text):
    response = client.embeddings.create(
        input=text,
        model="text-embedding-3-small",   # your deployment name
    )
    return response.data[0].embedding

Azure AI Search also supports a feature called integrated vectorization. With this feature, the service itself calls a Foundry Models embedding deployment to convert text into vectors at indexing time and at query time, so you do not have to generate the embeddings in your own code. This is the same mechanism Foundry IQ leans on under its managed layer — convenient for larger pipelines, but the basic flow shown above is enough to understand how vector search works underneath it.

A few practical notes

When you move from this small example to a real system, please keep the following points in mind. Choose your embedding model carefully, because it decides the dimension of your vectors and the quality of your retrieval. Add a chunking step so that long documents are split into focused passages. If you need both keyword matching and meaning-based matching, use hybrid search by passing a value to search_text together with the vector query; Azure AI Search will combine the keyword (BM25) results and the vector results. For even better ordering, you can enable the semantic ranker, which re-ranks the top results using a language model. Finally, monitor your index size, because vector indexes consume memory; if storage becomes a concern, look at vector compression and binary vectors. And if you'd rather not manage index freshness and chunking yourself, Foundry IQ wraps this same Azure AI Search foundation in a managed layer with automatic re-indexing — worth evaluating once your custom pipeline is working, to see whether the managed tradeoff is right for your document set.

Conclusion

We have seen what vector search is, why a RAG system needs it, and how Azure AI Search provides it through vector fields, the HNSW and exhaustive KNN algorithms, and the profile-based configuration. We then built a small but complete example: we created a search service, defined a vector index, embedded and uploaded a few documents, searched them by meaning, and assembled a RAG prompt. We also saw how to switch to Foundry Models embeddings.

You can now extend this example with your own documents, a chunking step, hybrid search, and a language model of your choice — deployed through Microsoft Foundry or otherwise — to build a full RAG application.

References

Quickstart: Vector search in Azure AI Search — Microsoft Learn: https://learn.microsoft.com/en-us/azure/search/search-get-started-vector
Create a vector index — Microsoft Learn: https://learn.microsoft.com/en-us/azure/search/vector-search-how-to-create-index
Vector search overview — Microsoft Learn: https://learn.microsoft.com/en-us/azure/search/vector-search-overview
Hybrid search overview — Microsoft Learn: https://learn.microsoft.com/en-us/azure/search/hybrid-search-overview
Index binary vectors (memory optimization) — Microsoft Learn: https://learn.microsoft.com/en-us/azure/search/vector-search-how-to-index-binary-data
azure-search-documents Python SDK — PyPI: https://pypi.org/project/azure-search-documents/
Python vector search sample — Azure SDK for Python (GitHub): https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/search/azure-search-documents/samples/sample_vector_search.py
Azure AI Search vector samples — GitHub: https://github.com/Azure/azure-search-vector-samples
Microsoft Foundry overview — Microsoft Learn: https://learn.microsoft.com/en-us/azure/ai-foundry/
Foundry IQ / knowledge bases overview — Microsoft Learn: https://learn.microsoft.com/en-us/azure/ai-foundry/concepts/knowledge-base

Amazon OpenSearch Vector Search Explained for RAG Systems

Jubin Soni — Tue, 02 Jun 2026 02:02:29 +0000

In this article, we will understand how vector search works in Amazon OpenSearch and how to use it as the retrieval layer in a Retrieval-Augmented Generation (RAG) system. The article is meant for software engineers. We will not stop at theory. We will build a small, working example that you can run on your own machine and follow along step by step.

By the end, you will have a small document search service that takes a user question, finds the most relevant text using vector similarity, and prepares the context that you can pass to a language model.

Let us begin.

What is a RAG system, in short

A RAG system has two main parts. The first part is retrieval. When a user asks a question, we search a knowledge base and pull out the most relevant pieces of text. The second part is generation. We pass these pieces of text, along with the question, to a language model so that the model can answer using real, grounded information.

Why we need vector search

Traditional keyword search matches exact words. If the user searches for "car" and the document says "automobile", a keyword search may miss it. Vector search solves this problem.

Amazon OpenSearch supports this through its k-NN (k-nearest neighbour) feature. We store embeddings in a special field type called knn_vector, and OpenSearch builds an index that can search through millions of vectors quickly.

How Amazon OpenSearch performs vector search

Amazon OpenSearch supports two broad approaches: exact nearest neighbour search and approximate nearest neighbour (ANN) search. Exact search compares the query against every stored vector and is accurate but slow on large data. ANN search trades a small amount of accuracy for a large gain in speed, which is what most production RAG systems use.

For ANN, Amazon OpenSearch provides different engines. An engine is the underlying library that builds and searches the vector index. The three engines are FAISS, Lucene, and NMSLIB. Each engine implements one or more algorithms. The two algorithms you will see most often are HNSW (Hierarchical Navigable Small World) and IVF (Inverted File).

The following table compares the three engines.

Engine	Algorithms supported	Filtering	Status	Best suited for
FAISS	HNSW and IVF	Efficient filtering (HNSW from version 2.9, IVF from 2.10)	Active and commonly recommended	Large data sets, memory tuning, and vector compression
Lucene	HNSW	Efficient filtering built in	Active	Smaller deployments and workloads that need metadata filtering
NMSLIB	HNSW	Pre-filtering and post-filtering only	Deprecated, kept for older indexes	Legacy indexes created in earlier versions

For new projects, FAISS and Lucene are the recommended choices. NMSLIB is deprecated, so please do not select it for new work. In our tutorial we will use the FAISS engine with the HNSW algorithm, because it is a good default for RAG workloads.

The knn_vector field supports up to 16,000 dimensions, and each dimension is stored as a 32-bit float by default. The number of dimensions is decided by the embedding model you choose. For example, the model we will use in this tutorial produces 384 dimensions.

The architecture of our RAG system

The data model

When we store a document for RAG, we usually do not store the full document as one record. We split it into smaller chunks, because a smaller chunk gives more focused retrieval and fits better inside the language model prompt. Each chunk becomes one record in the OpenSearch index, and each record holds the chunk text, its embedding, and some metadata for tracing the result back to its source.

The following entity relationship diagram shows how a source document relates to chunks, and how each chunk is stored as one vector record in OpenSearch.

Hands-on tutorial

Now we will build the system. I am assuming you have Docker and Python (version 3.9 or above) installed on your machine.

We will run OpenSearch locally using Docker. This is important to understand: Amazon OpenSearch Service is the managed version of OpenSearch, and the API is the same. So the code we write against a local OpenSearch will work against an Amazon OpenSearch Service domain with only a change in the connection settings. I will show you that change at the end.

Step 1: Run OpenSearch locally

Run the following command. It starts a single-node OpenSearch cluster. Recent OpenSearch versions require you to set an initial admin password, so we provide one through an environment variable.

docker run -d \
  --name opensearch-rag \
  -p 9200:9200 -p 9600:9600 \
  -e "discovery.type=single-node" \
  -e "OPENSEARCH_INITIAL_ADMIN_PASSWORD=Your-Strong-Password123!" \
  opensearchproject/opensearch:2.17.0

After a few seconds, check that it is running:

curl -k -u admin:'Your-Strong-Password123!' https://localhost:9200

You should receive a JSON response that includes the cluster name and version. The -k flag tells curl to accept the self-signed certificate, which is fine for local testing.

Step 2: Create the k-NN index

Now we create an index that has a knn_vector field. We must enable k-NN at the index level and define the method (the algorithm, the engine, and its parameters).

PUT /rag-docs
{
  "settings": {
    "index": {
      "knn": true,
      "knn.algo_param.ef_search": 100
    }
  },
  "mappings": {
    "properties": {
      "embedding": {
        "type": "knn_vector",
        "dimension": 384,
        "method": {
          "name": "hnsw",
          "space_type": "innerproduct",
          "engine": "faiss",
          "parameters": {
            "ef_construction": 128,
            "m": 16
          }
        }
      },
      "doc_id": { "type": "keyword" },
      "source": { "type": "keyword" },
      "text":   { "type": "text" }
    }
  }
}

Let me explain the choices here. The dimension is 384 because that is the output size of our embedding model. The space_type is innerproduct. We will produce normalised embeddings, and for unit-length vectors the inner product is equal to cosine similarity. This is a common and reliable pattern. The HNSW parameters m and ef_construction control the structure of the graph, and ef_search controls how thoroughly the graph is searched at query time.

The next table explains these parameters and their trade-offs.

Parameter	Where it is set	What it controls	Trade-off
`m`	Index mapping (method)	Number of connections each node keeps in the HNSW graph	Higher value gives better recall but uses more memory
`ef_construction`	Index mapping (method)	Number of candidate neighbours examined while building the graph	Higher value gives a better graph but slows down indexing
`ef_search`	Index settings	Number of candidate neighbours examined during a query	Higher value gives better recall but slows down each search
`space_type`	Index mapping (method)	The distance metric, such as `l2`, `innerproduct`, or `cosinesimil`	Must match the way your embeddings are produced

The default values (m is 16, ef_construction is around 128, ef_search is 100) are a reasonable starting point. You can tune them later based on your recall and latency requirements.

Step 3: Connect, embed, and index the documents

Install the required Python libraries:

pip install opensearch-py sentence-transformers

Now we write a Python script. First we connect to OpenSearch, then we load the embedding model, then we embed a few documents and store them.

from opensearchpy import OpenSearch, helpers
from sentence_transformers import SentenceTransformer

# 1. Connect to the local OpenSearch cluster
client = OpenSearch(
    hosts=[{"host": "localhost", "port": 9200}],
    http_auth=("admin", "Your-Strong-Password123!"),
    use_ssl=True,
    verify_certs=False,   # local self-signed certificate
)

# 2. Load the embedding model (produces 384-dimensional vectors)
model = SentenceTransformer("all-MiniLM-L6-v2")

# 3. Our small knowledge base
documents = [
    {"doc_id": "1", "source": "billing-faq",
     "text": "You can update your payment method from the account settings page under Billing."},
    {"doc_id": "2", "source": "billing-faq",
     "text": "Refunds are processed within five to seven business days to the original payment method."},
    {"doc_id": "3", "source": "shipping-faq",
     "text": "Standard delivery takes three to five working days within the country."},
    {"doc_id": "4", "source": "account-faq",
     "text": "To reset your password, click the forgot password link on the login screen."},
    {"doc_id": "5", "source": "shipping-faq",
     "text": "International orders may take up to fourteen working days depending on customs."},
]

# 4. Create embeddings. We normalise them so inner product equals cosine similarity.
texts = [d["text"] for d in documents]
embeddings = model.encode(texts, normalize_embeddings=True)

# 5. Build the bulk request and index everything
actions = []
for doc, vector in zip(documents, embeddings):
    actions.append({
        "_index": "rag-docs",
        "_source": {
            "doc_id": doc["doc_id"],
            "source": doc["source"],
            "text": doc["text"],
            "embedding": vector.tolist(),
        },
    })

helpers.bulk(client, actions)
client.indices.refresh(index="rag-docs")
print("Indexed", len(documents), "documents.")

When you run this script, it will index five small documents. In a real project you would read documents from files or a database, split them into chunks, and index thousands or millions of records using the same helpers.bulk approach.

Step 4: Search using a query vector

Now we write the retrieval function. We embed the user question with the same model, then we run a k-NN query. The query asks OpenSearch to return the k nearest vectors to our query vector.

def search(question, k=3):
    # Embed the question with the same model and normalisation
    query_vector = model.encode([question], normalize_embeddings=True)[0]

    body = {
        "size": k,
        "query": {
            "knn": {
                "embedding": {
                    "vector": query_vector.tolist(),
                    "k": k,
                }
            }
        },
    }

    response = client.search(index="rag-docs", body=body)
    results = []
    for hit in response["hits"]["hits"]:
        results.append({
            "score": hit["_score"],
            "source": hit["_source"]["source"],
            "text": hit["_source"]["text"],
        })
    return results


for item in search("how do I get my money back", k=3):
    print(round(item["score"], 4), "|", item["source"], "|", item["text"])

Step 5: Build the RAG prompt

The retrieval part is now complete. The final step is to take the retrieved text and build a prompt for the language model. We do not call any specific model here, because you may use Amazon Bedrock, an Anthropic model, an OpenAI model, or any other. We only prepare the input.

def build_prompt(question, k=3):
    hits = search(question, k=k)
    context = "\n\n".join(f"- {hit['text']}" for hit in hits)

    prompt = (
        "You are a support assistant. Use only the context below to answer "
        "the question. If the answer is not in the context, say that you do "
        "not have enough information.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {question}\n"
        "Answer:"
    )
    return prompt


print(build_prompt("how do I get my money back"))

The output is a prompt that contains the question and the most relevant pieces of text. You would now send this prompt to your language model, and the model would generate a grounded answer. This is the complete retrieval-augmented generation flow, with Amazon OpenSearch acting as the vector store and retriever.

Running this against Amazon OpenSearch Service

So far we have used a local cluster. To run the exact same code against an Amazon OpenSearch Service domain, you only need to change how the client connects. If your domain uses IAM-based access, you sign your requests with your AWS credentials. The rest of the code (creating the index, indexing, and searching) stays the same.

from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth
import boto3

host = "my-domain.us-east-1.es.amazonaws.com"   # your domain endpoint, no https://
region = "us-east-1"
service = "es"   # use "aoss" for Amazon OpenSearch Serverless

credentials = boto3.Session().get_credentials()
auth = AWSV4SignerAuth(credentials, region, service)

client = OpenSearch(
    hosts=[{"host": host, "port": 443}],
    http_auth=auth,
    use_ssl=True,
    verify_certs=True,
    connection_class=RequestsHttpConnection,
    pool_maxsize=20,
)

One point to remember: if you use Amazon OpenSearch Serverless vector search collections, the service value is aoss, and these collections support only the HNSW algorithm with the FAISS engine. They do not support the Lucene engine or the IVF algorithm. For a managed OpenSearch domain (not serverless), all three engines are available.

A few practical notes for production

Conclusion

We have seen what vector search is, why a RAG system needs it, and how Amazon OpenSearch provides it through the knn_vector field and the FAISS, Lucene, and NMSLIB engines. We then built a small but complete example: we ran OpenSearch, created a k-NN index, embedded and indexed a few documents, searched them by meaning, and assembled a RAG prompt. The same code runs against Amazon OpenSearch Service with only a change in the connection settings.

You can now extend this example with your own documents, a chunking step, and a language model of your choice to build a full RAG application.

References

Amazon OpenSearch Service vector database capabilities revisited — AWS Big Data Blog: https://aws.amazon.com/blogs/big-data/amazon-opensearch-service-vector-database-capabilities-revisited/
Approximate k-NN search — OpenSearch Documentation: https://docs.opensearch.org/latest/vector-search/vector-search-techniques/approximate-knn/
Methods and engines — OpenSearch Documentation: https://docs.opensearch.org/latest/mappings/supported-field-types/knn-methods-engines/
k-NN vector field type — OpenSearch Documentation: https://docs.opensearch.org/latest/mappings/supported-field-types/knn-vector/
Efficient filters in the OpenSearch vector engine — OpenSearch Blog: https://opensearch.org/blog/efficient-filters-in-knn/
Signing HTTP requests to Amazon OpenSearch Service (Python client) — AWS Documentation: https://docs.aws.amazon.com/opensearch-service/latest/developerguide/request-signing.html
Working with vector search collections (Amazon OpenSearch Serverless) — AWS Documentation: https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vector-search.html
Memory-optimized vectors — OpenSearch Documentation: https://docs.opensearch.org/latest/mappings/supported-field-types/knn-memory-optimized/

DEV Community: Jubin Soni

Implementing Zero-Trust Networking and Identity for Microsoft Foundry Agents

Zero-trust network boundary

Private endpoints: keeping traffic off the public internet

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

Data residency and customer-managed keys

Audit trails that actually hold up

Testing the security boundary, not just the happy path

References

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

Continuous evaluation pipeline

Evaluation: measuring quality continuously, not just at launch

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

Chargeback tagging: attributing cost to the right owner

Observability: the on-call-ready dashboard

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

References

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

Architecture Comparison

Decision Flow

Detailed Capability Comparison

Step 1 — Reading Data from Fabric OneLake in Azure Databricks

Step 2 — Writing Databricks Results Back to OneLake

Step 3 — When to Use Fabric Notebooks vs Databricks Notebooks

When to Use Which: Decision Framework

Things to Watch in Production

Wrapping Up

References

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

Architecture Overview

CI/CD Flow

Pipeline Stage Breakdown

Step 1 — Project Structure with Databricks Asset Bundles

Step 2 — Training Job with MLflow Autologging

Step 3 — Validation Job: Gate on Metrics Before Promoting

Step 4 — GitHub Actions CI/CD Workflow

Step 5 — Update the Serving Endpoint

Tool Comparison

Things to Watch in Production

Wrapping Up

References

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

Architecture Overview

Pipeline Flow

Layer Breakdown

Step 1 — Bronze Layer: Raw Ingest

Step 2 — Silver Layer: Clean and Validate

Step 3 — Gold Layer: Feature Engineering

Step 4 — MLflow: Track the Training Run

Bonus: Delta Lake Time Travel for Feature Reproducibility

Service Comparison

Things to Watch in Production

Wrapping Up

References

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

Architecture Overview

Event Flow

Code: Publishing Events to Kinesis

Code: Lambda Processor — Enrich and Invoke Bedrock

Code: Serving Recommendations via Lambda API

Service Comparison: Why Each AWS Service?

Things to Watch in Production

Wrapping Up

References

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

Table of Contents

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

Key Takeaways

References

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

Table of Contents

Why Query Optimization Matters