Github user WeichenXu123 commented on a diff in the pull request:
https://github.com/apache/spark/pull/19229#discussion_r140675967
--- Diff: mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala ---
@@ -223,20 +223,18 @@ class ImputerModel private[ml] (
override def transform(dataset: Dataset[_]): DataFrame = {
transformSchema(dataset.schema, logging = true)
- var outputDF = dataset
val surrogates = surrogateDF.select($(inputCols).map(col):
_*).head().toSeq
- $(inputCols).zip($(outputCols)).zip(surrogates).foreach {
+ val newCols = $(inputCols).zip($(outputCols)).zip(surrogates).map {
case ((inputCol, outputCol), surrogate) =>
val inputType = dataset.schema(inputCol).dataType
val ic = col(inputCol)
- outputDF = outputDF.withColumn(outputCol,
- when(ic.isNull, surrogate)
+ when(ic.isNull, surrogate)
--- End diff --
style: indent
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
