Ensemble pipeline component in Spark

Ensemble.scala

class Ensemble(val uid: String, models: Seq[RegressionModel[_, _]]) extends Model[RegressionModel[_, _]] {
  import org.apache.spark.sql.functions._
  
  def this(models: Seq[Model[_]]) = this(Identifiable.randomUID("ensemble"), models)

  override def copy(extra: ParamMap) = ???
  
  override def transform(
    dataset: Dataset[_]): DataFrame = {

    val predCols = models.map { m =>
      val preds = m.transform(dataset)
      preds.col(m.getPredictionCol)
    }
    
    def avgPreds = udf { row: Row =>
      val preds = row.toSeq.map(_.toString.toDouble)
      val size = preds.length
      preds.foldLeft(0.0)(_ + _) / size
    }
    dataset.withColumn("ensemble", avgPreds(predCols: _*))
  }

  override def transformSchema(
    schema: StructType): StructType = ???
}