Skip to content
Snippets Groups Projects
Commit c3ba4d4c authored by Reynold Xin's avatar Reynold Xin
Browse files

[MLlib] Minor: UDF style update.

Author: Reynold Xin <rxin@databricks.com>

Closes #4388 from rxin/mllib-style and squashes the following commits:

61d465b [Reynold Xin] oops
3364295 [Reynold Xin] Missed one ..
5e068e3 [Reynold Xin] [MLlib] Minor: UDF style update.
parent 7d789e11
No related branches found
No related tags found
No related merge requests found
......@@ -132,12 +132,14 @@ class LogisticRegressionModel private[ml] (
override def transform(dataset: DataFrame, paramMap: ParamMap): DataFrame = {
transformSchema(dataset.schema, paramMap, logging = true)
val map = this.paramMap ++ paramMap
val scoreFunction = udf((v: Vector) => {
val scoreFunction = udf { v: Vector =>
val margin = BLAS.dot(v, weights)
1.0 / (1.0 + math.exp(-margin))
} : Double)
}
val t = map(threshold)
val predictFunction = udf((score: Double) => { if (score > t) 1.0 else 0.0 } : Double)
val predictFunction = udf { score: Double =>
if (score > t) 1.0 else 0.0
}
dataset
.select($"*", scoreFunction(col(map(featuresCol))).as(map(scoreCol)))
.select($"*", predictFunction(col(map(scoreCol))).as(map(predictionCol)))
......
......@@ -129,13 +129,13 @@ class ALSModel private[ml] (
// Register a UDF for DataFrame, and then
// create a new column named map(predictionCol) by running the predict UDF.
val predict = udf((userFeatures: Seq[Float], itemFeatures: Seq[Float]) => {
val predict = udf { (userFeatures: Seq[Float], itemFeatures: Seq[Float]) =>
if (userFeatures != null && itemFeatures != null) {
blas.sdot(k, userFeatures.toArray, 1, itemFeatures.toArray, 1)
} else {
Float.NaN
}
} : Float)
}
dataset
.join(users, dataset(map(userCol)) === users("id"), "left")
.join(items, dataset(map(itemCol)) === items("id"), "left")
......
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