[SOLVED] How to map over DataFrame in spark to extract RowData and make predictions using h2o mojo model

How to map over DataFrame in spark to extract RowData and make predictions using h2o mojo model

I have a saved h2o model in mojo format, and now I am trying to load it and use it to make predictions on a new dataset (df) as part of a spark app written in scala. Ideally, I wish to append a new row to the existing DataFrame containing the class probability based on this model.

I can see how to apply a mojo to an individual row already in a RowData format (as per answer here), but I am not sure how to map over an existing DataFrame so that it is in the right format to make predictions using the mojo model. I have worked with DataFrames a fair bit, but never with the underlying RDDs.

Also, should this model be serialised / broadcast so that predictions can be done in parallel on a cluster, or will it be available to all executors as part of the map?

I have gotten this far:

# load mojo model and create easy predict model wrapper
val mojo = MojoModel.load("loca/path/to/mojo/mojo.zip")
val easyModel = new EasyPredictModelWrapper(mojo)

# map over spark DataFrame, converty to rdd, and make predictions on each row:
df.rdd.map { row =>
   val prediction = easyModel.predictBinomial(row).classProbabilities
   println(prediction)
   }

But my row variable is not in the right format for this to work. Any suggestions on what to try next?

EDIT: my DataFrame consists of 70 predictive feature columns which are a mixture of integers and category/factor columns. A very simple sample DataFrame:

val df = Seq(
  (0, 3, "cat1"),
  (1, 2, "cat2"),
  (2, 6, "cat1")
).toDF("id", "age", "category")

Solution

Use this function to prepare RowData object needed for H2O:

def rowToRowData(df: DataFrame, row: Row): RowData = {
  val rowAsMap = row.getValuesMap[Any](df.schema.fieldNames)
  val rowData = rowAsMap.foldLeft(new RowData()) { case (rd, (k,v)) => 
    if (v != null) { rd.put(k, v.toString) }
    rd
  }
  rowData
}