StringChecker.scala (MLeap side)

StringChecker.scala

package ml.combust.mleap.runtime.transformer.feature
 
import ml.combust.mleap.core.feature.StringCheckerModel
import ml.combust.mleap.core.types._
import ml.combust.mleap.runtime.frame.{FrameBuilder, Row, Transformer}
import ml.combust.mleap.runtime.function.{StructSelector, UserDefinedFunction}
 
import scala.util.Try
 
 
/**
  * Created by mostafam on 6/30/20.
  */
 
case class StringChecker(override val uid: String = Transformer.uniqueName("string_checker"),
                         override val shape: NodeShape,
                         override val model: StringCheckerModel) extends Transformer {
 
  val exec: UserDefinedFunction = UserDefinedFunction(
    (values: Row) => model(values.toSeq(0).asInstanceOf[String], values.toSeq(1).asInstanceOf[String]): Double,
    outputSchema.fields.head.dataType,
    Seq(SchemaSpec(inputSchema)))
 
  val outputCol: String = outputSchema.fields.head.name
  val inputCols: Seq[String] = inputSchema.fields.map(_.name)
  private val inputSelector: StructSelector = StructSelector(inputCols)
 
  override def transform[TB <: FrameBuilder[TB]](builder: TB): Try[TB] = {
    builder.withColumn(outputCol, inputSelector)(exec)
  }
}