Do not make 1000 AFTInput but just 2

apache · Jun 12, 2016 · 2c8adbf · 2c8adbf
1 parent c86ede8
commit 2c8adbf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
@@ -398,7 +398,7 @@ class AFTSurvivalRegressionSuite
     // should not throw an exception.
     val dataset = spark.createDataFrame(
       sc.parallelize(generateAFTInput(
-        1, Array(5.5), Array(0.8), 1000, 42, 1.0, 2.0, 2.0).take(2), numSlices = 3))
+        1, Array(5.5), Array(0.8), 2, 42, 1.0, 2.0, 2.0), numSlices = 3))
     val trainer = new AFTSurvivalRegression()
     trainer.fit(dataset)
   }