Skip to content
Snippets Groups Projects
Commit ec03866a authored by Dominik Dahlem's avatar Dominik Dahlem Committed by Xiangrui Meng
Browse files

[SPARK-11343][ML] Allow float and double prediction/label columns in RegressionEvaluator

mengxr, felixcheung

This pull request just relaxes the type of the prediction/label columns to be float and double. Internally, these columns are casted to double. The other evaluators might need to be changed also.

Author: Dominik Dahlem <dominik.dahlem@gmail.combination>

Closes #9296 from dahlem/ddahlem_regression_evaluator_double_predictions_27102015.
parent ecfb3e73
No related branches found
No related tags found
No related merge requests found
......@@ -23,7 +23,8 @@ import org.apache.spark.ml.param.shared.{HasLabelCol, HasPredictionCol}
import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
import org.apache.spark.mllib.evaluation.RegressionMetrics
import org.apache.spark.sql.{DataFrame, Row}
import org.apache.spark.sql.types.DoubleType
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{DoubleType, FloatType}
/**
* :: Experimental ::
......@@ -72,10 +73,13 @@ final class RegressionEvaluator @Since("1.4.0") (@Since("1.4.0") override val ui
@Since("1.4.0")
override def evaluate(dataset: DataFrame): Double = {
val schema = dataset.schema
SchemaUtils.checkColumnType(schema, $(predictionCol), DoubleType)
SchemaUtils.checkColumnType(schema, $(labelCol), DoubleType)
val predictionType = schema($(predictionCol)).dataType
require(predictionType == FloatType || predictionType == DoubleType)
val labelType = schema($(labelCol)).dataType
require(labelType == FloatType || labelType == DoubleType)
val predictionAndLabels = dataset.select($(predictionCol), $(labelCol))
val predictionAndLabels = dataset
.select(col($(predictionCol)).cast(DoubleType), col($(labelCol)).cast(DoubleType))
.map { case Row(prediction: Double, label: Double) =>
(prediction, label)
}
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment