Customer Lifetime Value (CLV) & Churn Prediction using Survival Analysis with PySpark

This project simulates a real-world MNC-scale Customer Lifetime Value (CLV) platform using:

• PySpark
• Spark MLlib AFT Survival Regression
• MLflow
• Airflow
• Snowflake
• S3 Data Lake
• Power BI/Tableau
• Survival Analysis (Right-Censored Data)

Business Problem

Traditional churn models fail because:

• Millions of customers are still active
• Actual churn date is unknown
• Data is right-censored

Install requirements

pip install pyspark
pip install mlflow
pip install pandas
pip install numpy
pip install snowflake-connector-python

Load Dataset into S3

aws s3 cp OnlineRetail.csv \
s3://clv-project/raw/

Create Spark Session

from pyspark.sql import SparkSession

spark = (
    SparkSession.builder
    .appName("CLV Survival Analysis")
    .getOrCreate()
)

Read Dataset

df = spark.read.csv(
    "s3://clv-project/raw/OnlineRetail.csv",
    header=True,
    inferSchema=True
)

df.show()

Generate Monetary Value

from pyspark.sql.functions import col

df = df.withColumn(
    "sales",
    col("Quantity") * col("UnitPrice")
)

Generate RFM Features

from pyspark.sql.functions import *

snapshot_date = df.select(
    max("InvoiceDate")
).collect()[0][0]

rfm = df.groupBy("CustomerID").agg(
    datediff(
        lit(snapshot_date),
        max("InvoiceDate")
    ).alias("recency"),

    countDistinct("InvoiceNo")
        .alias("frequency"),

    sum("sales")
        .alias("monetary")
)

Create Survival Labels

Assume for 90 days

rfm = rfm.withColumn(
    "churn_flag",
    when(col("recency") > 90, 1)
    .otherwise(0)
)

Duration Variable

rfm = rfm.withColumn(
    "duration",
    col("recency")
)
rfm = rfm.withColumn(
    "censor",
    when(col("churn_flag")==1,0)
    .otherwise(1)
)

Feature Engineering

from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(
    inputCols=[
        "recency",
        "frequency",
        "monetary"
    ],
    outputCol="features"
)

aft_df = assembler.transform(rfm)

Train AFT Survival Model

from pyspark.ml.regression import \
AFTSurvivalRegression

aft = AFTSurvivalRegression(
    featuresCol="features",
    labelCol="duration",
    censorCol="censor",
    predictionCol="prediction"
)
model = aft.fit(aft_df)

MLflow Tracking

import mlflow

mlflow.set_experiment(
    "Customer_Lifetime_Value"
)

with mlflow.start_run():

    model = aft.fit(aft_df)

    mlflow.log_param(
        "model",
        "AFTSurvivalRegression"
    )

    mlflow.log_metric(
        "scale",
        model.scale
    )

    mlflow.spark.log_model(
        model,
        "aft_model"
    )

Predict Customer Lifetime

predictions = model.transform(aft_df)

predictions.select(
    "CustomerID",
    "prediction"
).show()

CLV Calculation

predictions = predictions.withColumn(
    "monthly_revenue",
    col("monetary")/12
)

predictions = predictions.withColumn(
    "predicted_clv",
    col("prediction")
    * col("monthly_revenue")
    * 0.3
)

Generate Survival Curves

quantiles = [0.25,0.5,0.75]

aft = AFTSurvivalRegression(
    quantileProbabilities=quantiles,
    quantilesCol="quantiles"
)
survival_predictions = model.transform(aft_df)

Write Gold dataset

predictions.write.mode("overwrite")\
.parquet(
"s3://clv-project/gold/customer_clv"
)

Load into snowflake

sfOptions = {
    "sfURL":"xxxxx.snowflakecomputing.com",
    "sfUser":"user",
    "sfPassword":"password",
    "sfDatabase":"CLV_DB",
    "sfSchema":"PUBLIC",
    "sfWarehouse":"COMPUTE_WH"
}
predictions.write \
.format("snowflake") \
.options(**sfOptions) \
.option(
"dbtable",
"CUSTOMER_CLV"
) \
.mode("overwrite") \
.save()

Airflow DAG

from airflow import DAG
from airflow.operators.bash import BashOperator
from datetime import datetime

dag = DAG(
    "clv_pipeline",
    start_date=datetime(2025,1,1),
    schedule="@daily"
)

rfm = BashOperator(
    task_id="rfm_generation",
    bash_command=
    "spark-submit rfm_generator.py",
    dag=dag
)

train = BashOperator(
    task_id="train_model",
    bash_command=
    "spark-submit aft_training.py",
    dag=dag
)

predict = BashOperator(
    task_id="predict_clv",
    bash_command=
    "spark-submit predict_clv.py",
    dag=dag
)

rfm >> train >> predict

Production Optimization for 40M+ Users

df.repartition(500)
spark.conf.set(
"spark.sql.adaptive.enabled",
"true"
)

EMR Clustering

Core Nodes: 10 Task Nodes: 50+