Cómo calcular time-to-event (survival) en Apache Spark: paso a paso

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

Este tutorial muestra cómo calcular métricas de time-to-event (survival) en Apache Spark para analizar cuánto tiempo hasta un evento (por ejemplo churn o fallo) y cómo tratar la censura. Saber calcular time-to-event en Apache Spark es útil para monitorizar retención, mantenimiento predictivo y comparar grupos a gran escala.

Pré-requisitos

Python 3.x con PySpark instalado (Spark 3.x).
Acceso a un entorno local o cluster con recursos suficientes.
Conocimientos básicos de DataFrame en PySpark (select, groupBy, withColumn).

Paso 1: Entender los conceptos y montar un ejemplo simple

Time-to-event exige tres cosas: un identificador de entidad, tiempo de inicio/fin (o duración) y un indicador de evento (1=evento ocurrió, 0=censurado). Vamos a crear un ejemplo pequeño para simular usuarios con fechas de inicio y fecha del evento (o null si censurado).

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'])

Paso 2: Calcular duración y variable de censura

Convirtiendo strings a fechas, calculamos la duración en días entre start_date y event_date. Para filas censuradas (event_date null) usamos una fecha de corte (ej: fecha de análisis) y marcamos censura=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()

Paso 3: Agregar para obtener funciones de supervivencia por grupo

Una forma simple de estimar la curva de supervivencia es calcular, para cada duración t, cuántas entidades están en riesgo y cuántas tuvieron evento. Vamos a obtener una tabla de conteos por duración (estilo life-table).

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

# tabla por duración en días
life_table = (df2
    .groupBy('duration_days')
    .agg(
        count('*').alias('n_at_time'),
        _sum(col('event')).alias('events')
    )
    .orderBy('duration_days')
)

life_table.show()

Paso 4: Calcular supervivencia acumulada (Kaplan-Meier simplificado)

Kaplan-Meier calcula la supervivencia como producto de (1 - d_i / n_i) donde d_i son eventos y n_i número en riesgo en ese tiempo. En Spark vamos a usar operaciones de ventana para calcular el producto acumulado transformándolo en log para estabilidad numérica.

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

lt2 = (life_table
    .withColumn('at_risk', col('n_at_time'))  # en ejemplos simples asumimos n_at_time como riesgo en ese 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()

Paso 5: Ejemplo práctico por grupo (ej.: plan A vs plan B)

Si tiene un campo de grupo (por ejemplo plan), repetimos el proceso por grupo usando window particionado. Mostraré cómo agregar por grupo y calcular survival por grupo.

# acrescentar grupo de exemplo
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()

Verificar el resultado

Confirme que existen columnas duration_days y event en df2 y que life_table/lt_group muestran conteos no nulos. La columna survival debe comenzar en 1 (o próximo) y decrecer tras los eventos. Verifique con .show() y con conteos por grupo: df3.groupBy('plan').count().

Conclusión

Ahora tiene un flujo básico para calcular time-to-event en Apache Spark: preparar fechas, definir censura, agregar por tiempo y calcular una estimación de supervivencia (Kaplan-Meier simplificado). Próximos pasos: ajustar la definición de "n_at_risk" correctamente (considerando riesgos acumulados), tratar intervalos de tiempo amplios, y visualizar curvas en Python (matplotlib) después de recolectar resultados. Consejo: pruebe con datos reales y verifique casos de cero eventos para evitar divisiones por cero.