What Can Be Done with Databricks? A Comprehensive Use Case Guide

from pyspark.sql import SparkSession
from pyspark.sql.functions import *

# Read raw sales data from multiple sources
raw_sales = spark.read.format("csv") \
    .option("header", "true") \
    .load("s3://raw-data/sales/*.csv")

# Transform data: cleansing, enrichment, aggregation
transformed_sales = raw_sales \
    .filter(col("amount") > 0) \
    .withColumn("sale_date", to_date(col("timestamp"))) \
    .withColumn("revenue", col("amount") * col("quantity")) \
    .join(customer_dim, "customer_id") \
    .join(product_dim, "product_id")

# Write to Delta Lake for analytics
transformed_sales.write.format("delta") \
    .mode("overwrite") \
    .partitionBy("sale_date") \
    .save("/delta/sales_fact")

from delta.tables import *

# Read new customer data
new_customers = spark.read.format("json").load("s3://incoming/customers/")

# Load existing customer dimension
customer_table = DeltaTable.forPath(spark, "/delta/dim_customer")

# Merge with SCD Type 2 logic
customer_table.alias("target").merge(
    new_customers.alias("source"),
    "target.customer_id = source.customer_id"
).whenMatchedUpdate(
    condition = "target.is_current = true AND source.updated_at > target.updated_at",
    set = {
        "is_current": "false",
        "end_date": "source.updated_at"
    }
).whenNotMatchedInsert(
    values = {
        "customer_id": "source.customer_id",
        "name": "source.name",
        "email": "source.email",
        "start_date": "source.updated_at",
        "end_date": "null",
        "is_current": "true"
    }
).execute()

# JSON from APIs
json_data = spark.read.format("json").load("s3://data/events/*.json")

# Parquet from data lake
parquet_data = spark.read.format("parquet").load("s3://data/warehouse/*.parquet")

# CSV from legacy systems
csv_data = spark.read.format("csv") \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .load("s3://data/legacy/*.csv")

# XML from external vendors
xml_data = spark.read.format("xml") \
    .option("rowTag", "record") \
    .load("s3://data/external/*.xml")

# Unified processing
unified_df = json_data.union(parquet_data.union(csv_data))

# Read from Kafka streams (POS systems)
transactions = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "broker:9092") \
    .option("subscribe", "pos-transactions") \
    .load()

# Enrich with store and product dimensions
enriched = transactions \
    .join(stores_df, "store_id") \
    .join(products_df, "product_id") \
    .join(promotions_df, ["product_id", "date"], "left")

# Aggregate by region and product category
aggregated = enriched \
    .groupBy(
        window(col("timestamp"), "1 hour"),
        col("region"),
        col("category")
    ) \
    .agg(
        sum("amount").alias("total_revenue"),
        count("transaction_id").alias("transaction_count"),
        avg("amount").alias("avg_transaction_value")
    )

# Write to Delta Lake
aggregated.writeStream \
    .format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", "/delta/checkpoints/retail") \
    .start("/delta/retail_analytics")

-- Cohort analysis: customer retention over time
WITH first_purchase AS (
  SELECT
    customer_id,
    DATE_TRUNC('month', MIN(order_date)) as cohort_month
  FROM orders
  GROUP BY customer_id
),
monthly_activity AS (
  SELECT
    f.cohort_month,
    DATE_TRUNC('month', o.order_date) as activity_month,
    COUNT(DISTINCT o.customer_id) as active_customers
  FROM first_purchase f
  JOIN orders o ON f.customer_id = o.customer_id
  GROUP BY f.cohort_month, DATE_TRUNC('month', o.order_date)
)
SELECT
  cohort_month,
  activity_month,
  active_customers,
  active_customers / FIRST_VALUE(active_customers) OVER (
    PARTITION BY cohort_month
    ORDER BY activity_month
  ) * 100 as retention_rate
FROM monthly_activity
ORDER BY cohort_month, activity_month;

-- Identify top-performing products by region with trend analysis
SELECT
  region,
  product_name,
  month,
  revenue,
  RANK() OVER (PARTITION BY region, month ORDER BY revenue DESC) as rank,
  revenue - LAG(revenue) OVER (PARTITION BY region, product_name ORDER BY month) as mom_change,
  AVG(revenue) OVER (
    PARTITION BY region, product_name
    ORDER BY month
    ROWS BETWEEN 2 PRECEDING AND CURRENT ROW
  ) as rolling_3month_avg
FROM product_revenue
WHERE month >= '2024-01-01'
ORDER BY region, month, rank;

-- Risk exposure analysis with real-time position calculation
CREATE OR REPLACE TEMP VIEW risk_positions AS
SELECT
  trader_id,
  security_type,
  SUM(CASE WHEN trade_type = 'BUY' THEN quantity ELSE -quantity END) as net_position,
  SUM(CASE WHEN trade_type = 'BUY' THEN quantity * price ELSE -quantity * price END) as position_value,
  AVG(volatility) as avg_volatility,
  MAX(last_updated) as last_trade_time
FROM trades t
JOIN securities s ON t.security_id = s.security_id
WHERE trade_date >= CURRENT_DATE - INTERVAL 30 DAYS
GROUP BY trader_id, security_type;

-- Calculate Value at Risk (VaR) by trader
SELECT
  trader_id,
  SUM(position_value) as total_exposure,
  SUM(position_value * avg_volatility * 2.33) as var_99_percent,
  COUNT(DISTINCT security_type) as diversification_score,
  CASE
    WHEN SUM(position_value * avg_volatility * 2.33) > 1000000 THEN 'HIGH_RISK'
    WHEN SUM(position_value * avg_volatility * 2.33) > 500000 THEN 'MEDIUM_RISK'
    ELSE 'LOW_RISK'
  END as risk_category
FROM risk_positions
GROUP BY trader_id
HAVING SUM(ABS(position_value)) > 0
ORDER BY var_99_percent DESC;

from pyspark.ml.feature import VectorAssembler, StandardScaler
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml import Pipeline
import mlflow
import mlflow.spark

# Prepare features
feature_cols = ['tenure', 'monthly_charges', 'total_charges', 'num_products', 'support_calls']
assembler = VectorAssembler(inputCols=feature_cols, outputCol="features_raw")
scaler = StandardScaler(inputCol="features_raw", outputCol="features")

# Define model
rf = RandomForestClassifier(
    featuresCol="features",
    labelCol="churned",
    numTrees=100,
    maxDepth=10
)

# Create pipeline
pipeline = Pipeline(stages=[assembler, scaler, rf])

# MLflow experiment tracking
with mlflow.start_run(run_name="churn_prediction_v1"):
    # Train model
    model = pipeline.fit(train_df)

    # Make predictions
    predictions = model.transform(test_df)

    # Calculate metrics
    from pyspark.ml.evaluation import BinaryClassificationEvaluator
    evaluator = BinaryClassificationEvaluator(labelCol="churned", metricName="areaUnderROC")
    auc = evaluator.evaluate(predictions)

    # Log metrics and model
    mlflow.log_param("num_trees", 100)
    mlflow.log_param("max_depth", 10)
    mlflow.log_metric("auc", auc)
    mlflow.spark.log_model(model, "model")

    print(f"Model AUC: {auc:.4f}")

import tensorflow as tf
from tensorflow import keras
import mlflow.tensorflow

# Load preprocessed image data
train_images = spark.read.format("delta").load("/delta/product_images/train")
val_images = spark.read.format("delta").load("/delta/product_images/validation")

# Convert to TensorFlow datasets
train_dataset = tf.data.Dataset.from_tensor_slices((train_images.images, train_images.labels))
val_dataset = tf.data.Dataset.from_tensor_slices((val_images.images, val_images.labels))

# Define CNN model
model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(128, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dropout(0.5),
    keras.layers.Dense(10, activation='softmax')
])

# Compile and train
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# MLflow autologging for TensorFlow
mlflow.tensorflow.autolog()

with mlflow.start_run():
    history = model.fit(
        train_dataset.batch(32),
        validation_data=val_dataset.batch(32),
        epochs=50
    )

from databricks import automl

# Load customer data
df = spark.table("customers.churn_features")

# Run AutoML
summary = automl.classify(
    dataset=df,
    target_col="churned",
    primary_metric="f1",
    timeout_minutes=60,
    max_trials=100
)

# Best model is automatically logged to MLflow
print(f"Best model: {summary.best_trial.model_description}")
print(f"F1 Score: {summary.best_trial.metrics['f1_score']:.4f}")

# Generate notebook with best model code
summary.best_trial.to_notebook()

from databricks.feature_store import FeatureStoreClient

fs = FeatureStoreClient()

# Create feature table with patient history
patient_features = spark.sql("""
  SELECT
    patient_id,
    AVG(length_of_stay) as avg_los,
    COUNT(*) as num_admissions,
    SUM(CASE WHEN emergency_admission THEN 1 ELSE 0 END) as emergency_count,
    MAX(comorbidity_score) as max_comorbidity,
    DATEDIFF(CURRENT_DATE, MIN(admission_date)) as patient_history_days
  FROM admissions
  GROUP BY patient_id
""")

# Register features
fs.create_table(
    name="healthcare.patient_features",
    primary_keys=["patient_id"],
    df=patient_features,
    description="Aggregated patient admission history"
)

# Train readmission prediction model
training_set = fs.create_training_set(
    df=labels_df,
    feature_lookups=[
        FeatureLookup(
            table_name="healthcare.patient_features",
            lookup_key="patient_id"
        )
    ],
    label="readmitted_within_30days"
)

# AutoML handles feature engineering automatically
summary = automl.classify(
    dataset=training_set.load_df(),
    target_col="readmitted_within_30days",
    timeout_minutes=90
)

from pyspark.sql.functions import *

# Read from IoT Hub/Event Hub
sensor_stream = spark.readStream \
    .format("eventhubs") \
    .option("eventhubs.connectionString", connection_string) \
    .load()

# Parse and process sensor data
parsed_stream = sensor_stream \
    .selectExpr("CAST(body AS STRING) as json_data") \
    .select(from_json(col("json_data"), sensor_schema).alias("data")) \
    .select("data.*")

# Detect anomalies
anomalies = parsed_stream \
    .withColumn("temp_anomaly",
                when(col("temperature") > 100, True).otherwise(False)) \
    .withColumn("vibration_anomaly",
                when(col("vibration") > 50, True).otherwise(False)) \
    .filter(col("temp_anomaly") | col("vibration_anomaly"))

# Alert system
def send_alert(batch_df, batch_id):
    alerts = batch_df.collect()
    for alert in alerts:
        # Send to alerting system
        send_notification(alert)

anomalies.writeStream \
    .foreachBatch(send_alert) \
    .start()

# Aggregate metrics
metrics = parsed_stream \
    .groupBy(
        window(col("timestamp"), "5 minutes"),
        col("machine_id")
    ) \
    .agg(
        avg("temperature").alias("avg_temp"),
        max("temperature").alias("max_temp"),
        avg("vibration").alias("avg_vibration"),
        count("*").alias("reading_count")
    )

# Write to Delta Lake for analytics
metrics.writeStream \
    .format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", "/delta/checkpoints/iot") \
    .start("/delta/iot_metrics")

# Read clickstream from Kafka
clicks = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "broker:9092") \
    .option("subscribe", "website-clicks") \
    .load()

# Parse and sessionize
sessionized = clicks \
    .selectExpr("CAST(value AS STRING) as json") \
    .select(from_json(col("json"), click_schema).alias("data")) \
    .select("data.*") \
    .withWatermark("timestamp", "10 minutes") \
    .groupBy(
        window(col("timestamp"), "30 minutes"),
        col("user_id")
    ) \
    .agg(
        count("*").alias("page_views"),
        countDistinct("page_url").alias("unique_pages"),
        sum(col("time_on_page")).alias("total_time"),
        collect_list("page_url").alias("page_sequence")
    )

# Real-time personalization
def update_recommendations(batch_df, batch_id):
    # Update recommendation engine
    batch_df.write.format("delta").mode("append").save("/delta/user_sessions")

    # Trigger ML inference
    recommendations = recommendation_model.transform(batch_df)
    recommendations.write.format("redis").save()

sessionized.writeStream \
    .foreachBatch(update_recommendations) \
    .start()

from delta.tables import *

# Query historical data
df_yesterday = spark.read.format("delta") \
    .option("versionAsOf", 5) \
    .load("/delta/transactions")

# Time travel with timestamp
df_last_week = spark.read.format("delta") \
    .option("timestampAsOf", "2024-01-07") \
    .load("/delta/transactions")

# View change history
delta_table = DeltaTable.forPath(spark, "/delta/transactions")
history = delta_table.history()
history.select("version", "timestamp", "operation", "operationMetrics").show()

# Rollback to previous version
delta_table.restoreToVersion(5)

# Multiple users can safely update simultaneously
deltaTable = DeltaTable.forPath(spark, "/delta/inventory")

# User 1: Update stock levels
deltaTable.update(
    condition = "product_id = 'ABC123'",
    set = {"stock_level": "stock_level - 5"}
)

# User 2: Update pricing (concurrent operation)
deltaTable.update(
    condition = "category = 'Electronics'",
    set = {"price": "price * 0.9"}  # 10% discount
)

# User 3: Delete discontinued products (concurrent operation)
deltaTable.delete("discontinued = true")

# Original schema
original_data = spark.read.format("delta").load("/delta/customers")

# New data with additional columns
new_data = spark.read.format("json").load("s3://new-data/customers_v2/")

# Automatic schema merging
new_data.write.format("delta") \
    .mode("append") \
    .option("mergeSchema", "true") \
    .save("/delta/customers")

# Schema evolution with constraints
deltaTable = DeltaTable.forPath(spark, "/delta/customers")
deltaTable.alter().addColumns([
    "loyalty_points INT",
    "preferred_contact STRING"
])

# Cell 1: Python - Data Engineer prepares data
df = spark.read.format("delta").load("/delta/sales")
df.createOrReplaceTempView("sales_data")

-- Cell 2: SQL - Data Analyst queries data
SELECT
  product_category,
  SUM(revenue) as total_revenue,
  COUNT(DISTINCT customer_id) as unique_customers
FROM sales_data
WHERE sale_date >= '2024-01-01'
GROUP BY product_category
ORDER BY total_revenue DESC

# Cell 3: Python - Data Scientist builds model
results = spark.sql("SELECT * FROM sales_data WHERE revenue > 1000")
ml_features = results.select("customer_id", "revenue", "frequency", "recency")

from sklearn.cluster import KMeans
# Build customer segmentation model

# Cell 4: R - Statistician performs analysis
library(SparkR)
sales_r <- sql("SELECT * FROM sales_data")
correlation_matrix <- cor(sales_r[, c("revenue", "quantity", "discount")])

# Databricks CLI - Sync notebooks with Git
databricks repos create \
  --url https://github.com/company/data-science \
  --provider github

# Pull latest changes
databricks repos update --repo-id 123 --branch main

# Commit changes
databricks repos commit \
  --repo-id 123 \
  --message "Add customer segmentation notebook"

# Create job via API
import requests

job_config = {
    "name": "daily_etl_pipeline",
    "tasks": [
        {
            "task_key": "extract_data",
            "notebook_task": {
                "notebook_path": "/Shared/ETL/extract",
                "base_parameters": {"date": "{{run_date}}"}
            },
            "new_cluster": {
                "spark_version": "13.3.x-scala2.12",
                "node_type_id": "i3.xlarge",
                "num_workers": 4
            }
        },
        {
            "task_key": "transform_data",
            "depends_on": [{"task_key": "extract_data"}],
            "notebook_task": {
                "notebook_path": "/Shared/ETL/transform"
            },
            "existing_cluster_id": "0701-cluster-id"
        },
        {
            "task_key": "load_data",
            "depends_on": [{"task_key": "transform_data"}],
            "notebook_task": {
                "notebook_path": "/Shared/ETL/load"
            }
        }
    ],
    "schedule": {
        "quartz_cron_expression": "0 0 2 * * ?",
        "timezone_id": "America/New_York"
    },
    "email_notifications": {
        "on_failure": ["[email protected]"]
    }
}

response = requests.post(
    f"{workspace_url}/api/2.1/jobs/create",
    headers={"Authorization": f"Bearer {token}"},
    json=job_config
)

from airflow import DAG
from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from datetime import datetime, timedelta

default_args = {
    'owner': 'data-engineering',
    'depends_on_past': False,
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'databricks_ml_pipeline',
    default_args=default_args,
    schedule_interval='@daily',
    start_date=datetime(2024, 1, 1)
)

# Task 1: Data ingestion
ingest_data = DatabricksSubmitRunOperator(
    task_id='ingest_data',
    databricks_conn_id='databricks_default',
    existing_cluster_id='cluster-id',
    notebook_task={
        'notebook_path': '/Production/Ingestion/daily_ingest',
        'base_parameters': {'date': '{{ ds }}'}
    },
    dag=dag
)

# Task 2: Feature engineering
engineer_features = DatabricksSubmitRunOperator(
    task_id='engineer_features',
    databricks_conn_id='databricks_default',
    existing_cluster_id='cluster-id',
    notebook_task={
        'notebook_path': '/Production/Features/compute_features'
    },
    dag=dag
)

# Task 3: Model training
train_model = DatabricksSubmitRunOperator(
    task_id='train_model',
    databricks_conn_id='databricks_default',
    new_cluster={
        'spark_version': '13.3.x-scala2.12',
        'node_type_id': 'i3.xlarge',
        'num_workers': 8,
        'spark_conf': {'spark.speculation': 'true'}
    },
    notebook_task={
        'notebook_path': '/Production/ML/train'
    },
    dag=dag
)

# Task 4: Model deployment
deploy_model = DatabricksSubmitRunOperator(
    task_id='deploy_model',
    databricks_conn_id='databricks_default',
    existing_cluster_id='cluster-id',
    notebook_task={
        'notebook_path': '/Production/ML/deploy'
    },
    dag=dag
)

# Define dependencies
ingest_data >> engineer_features >> train_model >> deploy_model

# Real-time fraud detection on transactions
transactions = spark.readStream \
    .format("kafka") \
    .option("subscribe", "transactions") \
    .load()

# Feature engineering
features = transactions \
    .withColumn("hour_of_day", hour(col("timestamp"))) \
    .withColumn("day_of_week", dayofweek(col("timestamp"))) \
    .withColumn("amount_zscore",
                (col("amount") - avg("amount").over(user_window)) /
                stddev("amount").over(user_window))

# Load fraud detection model
fraud_model = mlflow.pyfunc.load_model("models:/fraud_detection/Production")

# Score transactions in real-time
scored = fraud_model.predict(features)
alerts = scored.filter(col("fraud_probability") > 0.8)

# Send to security team
alerts.writeStream.foreachBatch(send_fraud_alert).start()

-- Identify high-risk patients for preventive care
SELECT
  patient_id,
  age,
  num_chronic_conditions,
  num_emergency_visits_12mo,
  total_healthcare_cost_12mo,
  CASE
    WHEN num_emergency_visits_12mo >= 3
         AND num_chronic_conditions >= 2 THEN 'VERY_HIGH'
    WHEN num_emergency_visits_12mo >= 2
         OR num_chronic_conditions >= 3 THEN 'HIGH'
    WHEN total_healthcare_cost_12mo > 50000 THEN 'MEDIUM'
    ELSE 'LOW'
  END as risk_score
FROM patient_summary
WHERE is_active = true

from pyspark.ml.recommendation import ALS

# Load user-product interactions
interactions = spark.read.format("delta").load("/delta/interactions")

# Train ALS model
als = ALS(
    userCol="user_id",
    itemCol="product_id",
    ratingCol="rating",
    coldStartStrategy="drop"
)

model = als.fit(interactions)

# Generate recommendations for all users
user_recs = model.recommendForAllUsers(10)

# Deploy to real-time serving
user_recs.write.format("delta").mode("overwrite").save("/delta/recommendations")

# Analyze cell tower performance
tower_metrics = spark.read.format("delta").load("/delta/tower_metrics")

# Identify congestion patterns
congestion_analysis = tower_metrics \
    .groupBy("tower_id", "hour_of_day") \
    .agg(
        avg("active_connections").alias("avg_connections"),
        max("active_connections").alias("peak_connections"),
        avg("bandwidth_utilization").alias("avg_bandwidth"),
        sum("dropped_calls").alias("total_dropped_calls")
    ) \
    .withColumn("congestion_score",
                col("avg_bandwidth") * 0.4 +
                col("total_dropped_calls") * 0.6)

# Recommend capacity expansion
high_congestion = congestion_analysis.filter(col("congestion_score") > 80)

-- Create catalog structure
CREATE CATALOG production;
CREATE SCHEMA production.customer_data;

-- Grant permissions by role
GRANT SELECT ON SCHEMA production.customer_data TO `data_analysts`;
GRANT MODIFY ON SCHEMA production.customer_data TO `data_engineers`;
GRANT ALL PRIVILEGES ON SCHEMA production.customer_data TO `data_admins`;

-- Column-level security
CREATE TABLE production.customer_data.customers (
  customer_id STRING,
  name STRING,
  email STRING,
  ssn STRING,
  purchase_history ARRAY<STRING>
) WITH (
  row_filter = "region = current_user_region()",
  column_mask_ssn = "CASE WHEN is_member('pii_readers') THEN ssn ELSE 'XXX-XX-XXXX' END"
);

# Query lineage information
lineage = spark.sql("""
  SELECT
    source_table,
    target_table,
    transformation_type,
    last_updated,
    updated_by
  FROM system.lineage.table_lineage
  WHERE target_table = 'production.customer_data.customers'
""")

-- Use Databricks SQL with Photon
SELECT
  product_category,
  COUNT(*) as order_count,
  SUM(order_value) as total_revenue
FROM orders
WHERE order_date >= '2024-01-01'
GROUP BY product_category;

# Job cluster with autoscaling and spot instances
cluster_config = {
    "autoscale": {
        "min_workers": 2,
        "max_workers": 20
    },
    "aws_attributes": {
        "availability": "SPOT_WITH_FALLBACK",
        "spot_bid_price_percent": 100
    },
    "spark_conf": {
        "spark.databricks.delta.optimizeWrite.enabled": "true",
        "spark.databricks.delta.autoCompact.enabled": "true"
    }
}