How to calculate time-to-event (survival) in Apache Spark: step by step

João Barros 10 de November de 2026 5 min read

This tutorial shows how to calculate time-to-event (survival) metrics in Apache Spark to analyze how long until an event (for example churn or failure) and how to handle censoring. Knowing how to calculate time-to-event in Apache Spark is useful to monitor retention, predictive maintenance and compare groups at scale.

Prerequisites

Python 3.x with PySpark installed (Spark 3.x).
Access to a local environment or cluster with sufficient resources.
Basic knowledge of DataFrame in PySpark (select, groupBy, withColumn).

Step 1: Understand the concepts and build a simple example

Time-to-event requires three things: an entity identifier, start/end time (or duration) and an event indicator (1=event occurred, 0=censored). Let’s create a small example to simulate users with start dates and event date (or null if censored).

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, datediff, to_date

spark = SparkSession.builder.appName("time_to_event_example").getOrCreate()

data = [
    (1, '2022-01-01', '2022-03-15'),
    (2, '2022-02-10', None),
    (3, '2022-01-20', '2022-01-30'),
    (4, '2022-03-01', None)
]

df = spark.createDataFrame(data, ['id','start_date','event_date'])

Step 2: Calculate duration and censoring variable

Converting strings to dates, we calculate the duration in days between start_date and event_date. For censored rows (event_date null) we use a cutoff date (e.g. analysis date) and mark censoring=0.

from pyspark.sql.functions import lit

analysis_date = to_date(lit('2022-04-01'))

df2 = (df
    .withColumn('start_date', to_date(col('start_date')))
    .withColumn('event_date', to_date(col('event_date')))
    .withColumn('obs_end', when(col('event_date').isNotNull(), col('event_date')).otherwise(analysis_date))
    .withColumn('duration_days', datediff(col('obs_end'), col('start_date')))
    .withColumn('event', when(col('event_date').isNotNull(), lit(1)).otherwise(lit(0)))
)

df2.select('id','start_date','event_date','obs_end','duration_days','event').show()

Step 3: Aggregate to obtain survival functions by group

A simple way to estimate the survival curve is to compute, for each duration t, how many entities are at risk and how many had an event. Let's get a table of counts by duration (life-table style).

from pyspark.sql.functions import count, sum as _sum

# table by duration in days
life_table = (df2
    .groupBy('duration_days')
    .agg(
        count('*').alias('n_at_time'),
        _sum(col('event')).alias('events')
    )
    .orderBy('duration_days')
)

life_table.show()

Step 4: Calculate cumulative survival (simplified Kaplan-Meier)

Kaplan-Meier computes survival as the product of (1 - d_i / n_i) where d_i are events and n_i number at risk at that time. In Spark we will use window operations to calculate the cumulative product by transforming to log for numerical stability.

from pyspark.sql.window import Window
from pyspark.sql.functions import log, exp, coalesce

lt2 = (life_table
    .withColumn('at_risk', col('n_at_time'))  # in simple examples we assume n_at_time as risk at that t
    .withColumn('q', col('events') / col('at_risk'))
    .withColumn('log_one_minus_q', log(1 - col('q')))
)

w = Window.orderBy('duration_days').rowsBetween(Window.unboundedPreceding, 0)

km = (lt2
    .withColumn('cum_log', coalesce(_sum(col('log_one_minus_q')).over(w), lit(0)))
    .withColumn('survival', exp(col('cum_log')))
    .select('duration_days','n_at_time','events','survival')
)

km.show()

Step 5: Practical example by group (e.g. plan A vs plan B)

If you have a group field (for example plan), repeat the process per group using a partitioned window. I will show how to aggregate by group and calculate survival per group.

# add example group
from pyspark.sql.functions import when

df3 = df2.withColumn('plan', when(col('id') % 2 == 0, 'B').otherwise('A'))

lt_group = (df3
    .groupBy('plan','duration_days')
    .agg(count('*').alias('n_at_time'), _sum(col('event')).alias('events'))
    .orderBy('plan','duration_days')
)

w_grp = Window.partitionBy('plan').orderBy('duration_days').rowsBetween(Window.unboundedPreceding, 0)

ltg2 = (lt_group
    .withColumn('q', col('events')/col('n_at_time'))
    .withColumn('log_one_minus_q', log(1 - col('q')))
    .withColumn('cum_log', coalesce(_sum(col('log_one_minus_q')).over(w_grp), lit(0)))
    .withColumn('survival', exp(col('cum_log')))
    .select('plan','duration_days','n_at_time','events','survival')
)

ltg2.show()

Verify the result

Confirm that there are columns duration_days and event in df2 and that life_table/lt_group show non-zero counts. The survival column should start at 1 (or close) and decrease after events. Check with .show() and with counts by group: df3.groupBy('plan').count().

Conclusion

You now have a basic pipeline to calculate time-to-event in Apache Spark: prepare dates, define censoring, aggregate by time and compute a survival estimate (simplified Kaplan-Meier). Next steps: adjust the definition of "n_at_risk" correctly (considering accumulated risk), handle large time intervals, and visualize curves in Python (matplotlib) after collecting results. Tip: test with real data and check for zero-event cases to avoid divisions by zero.