Train CoxPH Model in Sparkling Water¶

Sparkling Water provides API for H2O CoxPH in Scala and Python. The following sections describe how to train the CoxPH model in Sparkling Water in both languages. See also Parameters of H2OCoxPH.

Scala
Python

First, let’s start Sparkling Shell as

./bin/sparkling-shell

Start H2O cluster inside the Spark environment

import ai.h2o.sparkling._
import java.net.URI
val hc = H2OContext.getOrCreate()

Parse the data using H2O and convert them to Spark Frame

import org.apache.spark.SparkFiles
spark.sparkContext.addFile("https://raw.githubusercontent.com/h2oai/sparkling-water/master/examples/smalldata/coxph_test/heart.csv")
val heartDF = spark.read.option("header", "true").option("inferSchema", "true").csv(SparkFiles.get("heart.csv"))
val Array(trainingDF, testingDF) = heartDF.randomSplit(Array(0.8, 0.2), seed = 12345)

Train the model. You can configure all the available CoxPH parameters using provided setters.

import ai.h2o.sparkling.ml.algos.H2OCoxPH
val estimator = new H2OCoxPH().
    setStartCol("start").
    setStopCol("stop").
    setTies("breslow").
    setLabelCol("event")
val model = estimator.fit(trainingDF)

You can also get raw model details by calling the getModelDetails() method available on the model as:

model.getModelDetails()

Run Predictions

model.transform(testingDF).show(false)