[MINOR] Correct validateAndTransformSchema in GaussianMixture and AFTSurvivalRegression

sharp-pixel · srowen · commit 73e04ecc4f29 · 2017-08-20T11:05:54.000+01:00
## What changes were proposed in this pull request? The line SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType) did not modify the variable schema, hence only the last line had any effect. A temporary variable is used to correctly append the two columns predictionCol and probabilityCol. ## How was this patch tested? Manually. Please review http://spark.apache.org/contributing.html before opening a pull request. Author: Cédric Pelvet <cedric.pelvet@gmail.com> Closes apache#18980 from sharp-pixel/master.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala
@@ -64,8 +64,8 @@ private[clustering] trait GaussianMixtureParams extends Params with HasMaxIter w
    */
   protected def validateAndTransformSchema(schema: StructType): StructType = {
     SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT)
-    SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType)
-    SchemaUtils.appendColumn(schema, $(probabilityCol), new VectorUDT)
+    val schemaWithPredictionCol = SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType)
+    SchemaUtils.appendColumn(schemaWithPredictionCol, $(probabilityCol), new VectorUDT)
   }
 }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
@@ -109,10 +109,12 @@ private[regression] trait AFTSurvivalRegressionParams extends Params
       SchemaUtils.checkNumericType(schema, $(censorCol))
       SchemaUtils.checkNumericType(schema, $(labelCol))
     }
-    if (hasQuantilesCol) {
+
+    val schemaWithQuantilesCol = if (hasQuantilesCol) {
       SchemaUtils.appendColumn(schema, $(quantilesCol), new VectorUDT)
-    }
-    SchemaUtils.appendColumn(schema, $(predictionCol), DoubleType)
+    } else schema
+
+    SchemaUtils.appendColumn(schemaWithQuantilesCol, $(predictionCol), DoubleType)
   }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -64,8 +64,8 @@ private[clustering] trait GaussianMixtureParams extends Params with HasMaxIter w`
`64`	`64`	`*/`
`65`	`65`	`protected def validateAndTransformSchema(schema: StructType): StructType = {`
`66`	`66`	`SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT)`
`67`		`- SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType)`
`68`		`- SchemaUtils.appendColumn(schema, $(probabilityCol), new VectorUDT)`
	`67`	`+ val schemaWithPredictionCol = SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType)`
	`68`	`+ SchemaUtils.appendColumn(schemaWithPredictionCol, $(probabilityCol), new VectorUDT)`
`69`	`69`	`}`
`70`	`70`	`}`
`71`	`71`
Original file line number	Diff line number	Diff line change
`@@ -109,10 +109,12 @@ private[regression] trait AFTSurvivalRegressionParams extends Params`
`109`	`109`	`SchemaUtils.checkNumericType(schema, $(censorCol))`
`110`	`110`	`SchemaUtils.checkNumericType(schema, $(labelCol))`
`111`	`111`	`}`
`112`		`- if (hasQuantilesCol) {`
	`112`	`+`
	`113`	`+ val schemaWithQuantilesCol = if (hasQuantilesCol) {`
`113`	`114`	`SchemaUtils.appendColumn(schema, $(quantilesCol), new VectorUDT)`
`114`		`- }`
`115`		`- SchemaUtils.appendColumn(schema, $(predictionCol), DoubleType)`
	`115`	`+ } else schema`
	`116`	`+`
	`117`	`+ SchemaUtils.appendColumn(schemaWithQuantilesCol, $(predictionCol), DoubleType)`
`116`	`118`	`}`
`117`	`119`	`}`
`118`	`120`