Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SPARK-19234][MLLib] AFTSurvivalRegression should fail fast when any labels are zero #16652

Closed
wants to merge 6 commits into from
Original file line number Diff line number Diff line change
Expand Up @@ -644,4 +644,5 @@ private class AFTCostFun(
*/
private[regression] case class AFTPoint(features: Vector, label: Double, censor: Double) {
require(censor == 1.0 || censor == 0.0, "censor of class AFTPoint must be 1.0 or 0.0")
require(label > 0.0, "label of AFTPoint must be positive")
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,15 @@ package org.apache.spark.ml.regression

import scala.util.Random

import org.apache.spark.SparkFunSuite
import org.apache.spark.{SparkException, SparkFunSuite}
import org.apache.spark.ml.{Estimator, Model}
import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.ml.param.ParamsSuite
import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
import org.apache.spark.ml.util.TestingUtils._
import org.apache.spark.mllib.random.{ExponentialGenerator, WeibullGenerator}
import org.apache.spark.mllib.util.MLlibTestSparkContext
import org.apache.spark.sql.{DataFrame, Row}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

class AFTSurvivalRegressionSuite
extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
Expand Down Expand Up @@ -400,6 +401,19 @@ class AFTSurvivalRegressionSuite
val trainer = new AFTSurvivalRegression()
trainer.fit(dataset)
}

test("SPARK-19234: Fail fast on zero-valued labels") {
val dataset = spark.createDataFrame(Seq(
(1.218, 1.0, Vectors.dense(1.560, -0.605)),
(0.000, 0.0, Vectors.dense(0.346, 2.158)), // ← generates error; zero labels invalid
(4.199, 0.0, Vectors.dense(0.795, -0.226)))).toDF("label", "censor", "features")
val aft = new AFTSurvivalRegression()
withClue("label of AFTPoint must be positive") {
intercept[SparkException] {
aft.fit(dataset)
}
}
}
}

object AFTSurvivalRegressionSuite {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ object MLTestingUtils extends SparkFunSuite {
actuals.foreach(actual => check(expected, actual))

val dfWithStringLabels = spark.createDataFrame(Seq(
("0", 1, Vectors.dense(0, 2, 3), 0.0)
("1", 1, Vectors.dense(0, 2, 3), 0.0)
)).toDF("label", "weight", "features", "censor")
val thrown = intercept[IllegalArgumentException] {
estimator.fit(dfWithStringLabels)
Expand Down Expand Up @@ -156,7 +156,6 @@ object MLTestingUtils extends SparkFunSuite {
featuresColName: String = "features",
censorColName: String = "censor"): Map[NumericType, DataFrame] = {
val df = spark.createDataFrame(Seq(
(0, Vectors.dense(0)),
(1, Vectors.dense(1)),
(2, Vectors.dense(2)),
(3, Vectors.dense(3)),
Expand Down