File tree Expand file tree Collapse file tree 1 file changed +7
-11
lines changed
mllib/src/main/scala/org/apache/spark/mllib/classification Expand file tree Collapse file tree 1 file changed +7
-11
lines changed Original file line number Diff line number Diff line change @@ -266,12 +266,10 @@ class NaiveBayes private (
266
266
* (default: Multinomial)
267
267
*/
268
268
def setModelType (modelType: String ): NaiveBayes = {
269
- if (NaiveBayes .supportedModelTypes.contains(modelType)) {
270
- this .modelType = modelType
271
- this
272
- } else {
273
- throw new UnknownError (s " NaiveBayesModel does not support ModelType: $modelType" )
274
- }
269
+ require(NaiveBayes .supportedModelTypes.contains(modelType),
270
+ s " NaiveBayes was created with an unknown ModelType: $modelType" )
271
+ this .modelType = modelType
272
+ this
275
273
}
276
274
277
275
/** Get the model type. */
@@ -404,11 +402,9 @@ object NaiveBayes {
404
402
* multinomial or bernoulli
405
403
*/
406
404
def train (input : RDD [LabeledPoint ], lambda : Double , modelType : String ): NaiveBayesModel = {
407
- if (supportedModelTypes.contains(modelType)) {
408
- new NaiveBayes (lambda, modelType).run(input)
409
- } else {
410
- throw new UnknownError (s " NaiveBayes was created with an unknown ModelType: $modelType" )
411
- }
405
+ require(supportedModelTypes.contains(modelType),
406
+ s " NaiveBayes was created with an unknown ModelType: $modelType" )
407
+ new NaiveBayes (lambda, modelType).run(input)
412
408
}
413
409
414
410
}
You can’t perform that action at this time.
0 commit comments