Churn Prediction on Databricks:Governed ML with Unity Catalog

What sits here

• Gold-layer Delta tables — conformed customer, transaction and interaction facts
• Registered in Unity Catalog under a three-level namespace: catalog.schema.table
• Produced by the upstream ingestion pipelines — this architecture consumes them, it doesn't rebuild them

Why Delta specifically

• Time travel — training data can be reproduced exactly as it stood on the training date
• ACID transactions — no partially-written table feeding a training run
• Schema enforcement — an upstream change surfaces as an error, not as silently corrupted features

What gets pinned down

• Churn event — the precise, queryable definition (contract end, N days inactive, downgrade below threshold)
• Observation window — how much history feeds each example
• Prediction horizon — how far forward the label looks
• Intervention lead time — the horizon must exceed the time the business needs to act

Traps this avoids

• Label leakage — features computed after the churn event (a cancellation ticket, a final invoice) make training scores look excellent and production performance collapse
• Class imbalance — churn is usually rare; accuracy is a misleading metric and PR-AUC is reported instead
• Silent redefinition — the definition lives in version-controlled code, not in someone's head

Why Feature Views

• Features are declared as Unity Catalog objects — Databricks creates and manages the underlying pipelines
• Supports time-windowed aggregations natively, which is most of what churn features are
• Governed, discoverable and shareable across workspaces like any other UC asset
• Models trained on them automatically capture lineage back to the features used

Feature Views vs Feature Tables

• Feature Views — declarative, Databricks-managed pipelines; the recommended default
• Feature Tables — a UC Delta table with a primary key that you write to and own; used where the computation is too bespoke to declare

The mechanism

• Training sets are assembled with a point-in-time join: feature values as of the timestamp each label observation was recorded
• This is native to the Feature Store rather than hand-rolled with window functions
• It eliminates train/serve skew — the same lookup logic applies at training and at inference

Why hand-rolling this goes wrong

• Correct temporal joins across many feature tables at differing grains are genuinely difficult to get right
• Errors are silent — they produce optimistic offline metrics rather than exceptions
• The gap only surfaces in production, by which point the model is already distrusted

What it provides

• Feature values published to the Online Feature Store for low-latency retrieval
• Automatic feature lookup at inference — the serving endpoint fetches features by entity key
• On-demand computation for features that must be derived from request-time input

Why it's the same definition

• Offline training and online serving resolve from one feature definition, which is what keeps skew out
• A separately reimplemented "real-time version" of a feature is a well-known source of silent divergence

Captured per run

• Hyperparameters, metrics and artifacts
• Source notebook or bundle revision
• Data version — the Delta table version used, making the run reproducible via time travel
• Feature lineage, inherited automatically from Feature Store usage

MLflow 3 on Databricks

• Models capture parameters and metrics directly, surfaced on the model version page
• Runs and models are visible across workspaces, not siloed per workspace
• The registry URI defaults to databricks-uc, so UC is the registry by default

Role in the workflow

• Establishes a defensible baseline metric before bespoke development begins
• Glass-box — generates the underlying notebooks, so the baseline is inspectable and forkable rather than opaque
• Handles the standard preprocessing and model search, surfacing the leaderboard in MLflow

Why the baseline discipline matters

• Custom modelling effort has to be justified by measured lift over the baseline, not by assumption
• A weak baseline is diagnostic — it usually indicates a feature or label problem, and points effort at the right layer
• Generated notebooks are a legitimate starting point for the bespoke model rather than throwaway output

Metrics that fit the problem

• PR-AUC as the headline — appropriate under class imbalance where ROC-AUC flatters
• Precision@k — the operationally meaningful measure, since the team can only work a fixed-size list
• Calibration — predicted probabilities must be trustworthy if they drive prioritisation or spend
• Slice analysis — performance by segment, tenure and value band, to catch aggregate metrics hiding segment failure

Model signature

• Unity Catalog requires a model signature on every registered version — input and output schema are explicit
• This turns a schema mismatch into a clear error at deployment rather than corrupted scores in production

Three-level namespace

• Models are addressed as catalog.schema.model — e.g. prod.ml_churn.churn_model
• The catalog expresses the environment, which is how dev/staging/prod separation is represented
• Models sit in the same governance plane as tables and features — one permission model, not two

Privileges

What this replaces

• The workspace-scoped model registry, which was siloed per workspace and separate from data governance
• Lineage now spans data → features → model → predictions in a single graph

Aliases, not stages

• Unity Catalog does not support MLflow stages. Deployment is expressed through aliases — mutable named pointers to a specific version
• Loading by alias: models:/prod.ml_churn.churn_model@Champion
• Promotion is a pointer reassignment — no code change, no redeploy, instantly reversible

The pattern in practice

• @Champion serves production traffic
• @Challenger scores in parallel, writing to the same inference table with a model-version tag
• Comparison happens on realised outcomes, not offline metrics — the only comparison that settles the question
• Rollback is reassigning @Champion to the prior version

Deploy code, not artifacts

• Databricks recommends deploying ML pipelines as code — the training pipeline runs in the target environment and produces the model there
• This removes a whole class of environment-mismatch bugs, since the model is built where it will run
• Environments are separate catalogs, so promotion is expressed in the namespace

Where artifact promotion is still needed

• copy_model_version() copies a version across catalogs where retraining in prod isn't viable
• Requires read on source and write on destination — the permission boundary is the control
• The alias is then set on the destination model to activate it

The job

• Scheduled Databricks Workflow, loading the model by alias: models:/prod.ml_churn.churn_model@Champion
• Features resolved through the Feature Store, so scoring uses the same definitions as training
• Output written to a governed Delta table with model version and scoring timestamp stamped on every row

Explanations alongside scores

• Per-prediction feature attributions computed and stored with the score
• Adoption depends on this — a ranked list without reasons gets abandoned

The endpoint

• Serverless model serving, backed by the alias rather than a pinned version
• Inference tables enabled — requests and responses land automatically in a governed Delta table
• Online feature lookup by entity key, plus on-demand computation for request-time features

Why inference tables matter more than the latency

• They are the substrate the entire monitoring loop is built on — inputs, predictions and later ground truth in one governed place
• Without them, production behaviour has to be reconstructed from logs after the fact, if at all

Delivery paths

• Scores written back to the CRM so they appear in the existing workflow, not a parallel one
• AI/BI dashboards over the scored Delta table for management view
• Threshold-based alerting for high-value at-risk accounts

The feedback loop

• Intervention outcomes captured — what action was taken, and what followed
• Those outcomes become future labels, so the model improves against reality rather than a frozen snapshot
• Enables measurement of incremental retention lift, which is the number that justifies the programme

What sits under UC governance

• Delta tables across every medallion layer
• Feature Views and feature tables
• Registered models and versions
• Inference tables and scored output

Capabilities used

• Three-level namespace — catalog.schema.object, with catalog expressing environment
• Fine-grained access control — row filters and column masks applied at query time
• End-to-end lineage — table → feature → model → prediction as one graph
• Audit logs across data and model access alike

Monitor configuration

• Inference profile over the inference / scored table — the profile type built for model monitoring
• Tracks model inputs, predictions and, once available, ground-truth labels
• A baseline table pins the training distribution as the reference point

What it produces

• A profile metrics table — summary statistics per window
• A drift metrics table — statistical change versus baseline or the prior window
• An auto-generated dashboard, plus alerts driven off the metrics tables

The retraining loop

• Feature drift or a performance drop past threshold raises an alert and can trigger the training Workflow
• The candidate registers as a new version and takes the @Challenger alias
• Promotion stays a human decision — automated retraining, deliberately not automated deployment

What the bundle declares

• Workflow definitions — training, batch scoring, monitoring refresh
• Notebook and source paths, cluster and serverless configuration
• Per-environment target variables — catalog names, schedules, endpoint sizing

The pipeline

• Git-backed repo; every change is a branch and a pull request
• CI validates the bundle and runs unit tests on feature and training logic
• bundle deploy per environment target, driven from Azure DevOps — consistent with the existing release process
• Environment differences live in bundle variables, never in edited copies of a notebook

Compute strategy

• Job compute, not all-purpose, for scheduled work — cheaper per unit and terminates on completion
• Serverless where startup latency dominates runtime, avoiding idle spend entirely
• SQL warehouses with auto-stop for the BI layer
• Scoring frequency matched to how quickly churn risk actually moves — daily, not hourly

Where cost is really decided

• Batch versus real-time is the dominant lever — an always-on endpoint costs continuously, a nightly job costs for minutes
• Feature computation reused across training and scoring rather than recomputed per path
• Retraining triggered by drift rather than by a fixed aggressive schedule

Optimised against real failure modes

• Adoption — write-back into existing workflows, with per-prediction explanations; a standalone dashboard is an adoption risk, not a deliverable
• Silent decay — inference profiling and drift metrics make degradation visible rather than assumed
• Explainability — UC lineage spans data → feature → model → prediction as one graph
• Bus factor — Asset Bundles make the system reproducible from source
• Problem definition — label and horizon pinned in code before modelling starts

Deliberate choices worth defending

• Batch-first — real-time serving is opt-in, justified by a latency requirement rather than assumed
• Automated retraining, manual promotion — automating the expensive, repeatable half while keeping a human on the risky half
• Champion/Challenger on realised outcomes — offline metrics inform, live outcomes decide
• PR-AUC and precision@k — metrics matched to imbalance and to a fixed-capacity worklist

Trade-offs, honestly

• Heavier upfront investment than a notebook that produces a scored CSV — and for a genuine one-off analysis, the notebook is the right answer
• Feature Views constrain how features are expressed in exchange for managed pipelines and lineage
• Holdout groups for measuring true lift cost some retention in the short term to prove the programme works at all
• The monitoring loop only fully closes when ground truth arrives, which for churn is months later

Entry points

• Scheduled — periodic refresh aligned to how fast the underlying behaviour moves
• Drift-triggered — fired from the monitoring alert at step 13

Run context captured

• Trigger reason is logged on the run, so a model version can always be traced to why it was built

Recomputed per run

• Churn events resolved over the observation window using the version-controlled definition
• Each row carries its observation timestamp, which step 3 depends on
• Class balance logged — a sudden shift is itself a signal worth alerting on

Mechanism

• Feature Store training set built with a point-in-time join on the observation timestamp
• Feature lineage is captured automatically and travels with the model
• The identical lookup logic is reused at scoring time, eliminating train/serve skew

Logged to MLflow

• Hyperparameters, metrics, artifacts
• Delta table version — reproducible via time travel
• Model signature — required for Unity Catalog registration
• Source revision from the deployed bundle

Comparison

• Both models scored on a common holdout — PR-AUC and precision@k as primaries
• Calibration checked, since scores drive prioritisation
• Slice analysis by tenure and value band — an aggregate gain hiding a segment regression is a fail

Branch logic

• Pass — register a new version in UC, assign the @Challenger alias
• Fail — halt, leave @Champion untouched, raise an alert with the comparison attached
• Thresholds are explicit and version-controlled, not implicit in a notebook

The action

• Reviewer inspects the comparison, then reassigns @Champion to the new version
• No redeploy — serving and scoring resolve by alias, so the switch is instant
• Rollback is reassigning the alias back; the prior version is untouched

Why alias, not version

• The job references @Champion, so promotion and rollback require no code change or redeploy
• The resolved version number is logged on the run and stamped onto every scored row

Output

• Features resolved through the Feature Store — same definitions as training
• Scores plus SHAP attributions written to a governed Delta table
• Every row stamped with model version + scoring timestamp — the basis for later comparison

Publication paths

• CRM write-back so scores appear inside the existing workflow
• AI/BI dashboards over the scored Delta table
• Threshold alerting for high-value at-risk accounts
• Intervention outcomes captured back, becoming labels for a future step 2

What runs

• Monitor with an inference profile over the scored table
• Tracks model inputs, predictions, and ground-truth labels once they arrive
• Emits a profile metrics table per window, plus an auto-generated dashboard

Computation

• Drift metrics table — current window vs the pinned training baseline and vs the prior window
• Input drift is the leading indicator; outcome-based performance confirms it later

The closing edge

• Alerts defined over the drift and profile metrics tables
• Breach raises a notification and can fire the training Workflow at step 1
• Thresholds are explicit and reviewed — too tight retrains constantly, too loose defeats the purpose