Skip to content

Commit

Permalink
[MINOR] Correct validateAndTransformSchema in GaussianMixture and AFT…
Browse files Browse the repository at this point in the history
…SurvivalRegression

## 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 #18980 from sharp-pixel/master.

(cherry picked from commit 73e04ec)
Signed-off-by: Sean Owen <sowen@cloudera.com>
  • Loading branch information
sharp-pixel authored and srowen committed Aug 20, 2017
1 parent 6f366fb commit 2394ae2
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 5 deletions.
Expand Up @@ -62,8 +62,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)
}
}

Expand Down
Expand Up @@ -109,10 +109,12 @@ private[regression] trait AFTSurvivalRegressionParams extends Params
SchemaUtils.checkColumnType(schema, $(censorCol), DoubleType)
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)
}
}

Expand Down

0 comments on commit 2394ae2

Please sign in to comment.