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

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()}")

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()}")

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")

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}")

# 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