aboutsummaryrefslogtreecommitdiff
path: root/R
diff options
context:
space:
mode:
Diffstat (limited to 'R')
-rw-r--r--R/pkg/inst/tests/testthat/test_mllib.R10
1 files changed, 10 insertions, 0 deletions
diff --git a/R/pkg/inst/tests/testthat/test_mllib.R b/R/pkg/inst/tests/testthat/test_mllib.R
index a1eaaf2091..c99315726a 100644
--- a/R/pkg/inst/tests/testthat/test_mllib.R
+++ b/R/pkg/inst/tests/testthat/test_mllib.R
@@ -481,6 +481,16 @@ test_that("spark.naiveBayes", {
expect_error(m <- e1071::naiveBayes(Survived ~ ., data = t1), NA)
expect_equal(as.character(predict(m, t1[1, ])), "Yes")
}
+
+ # Test numeric response variable
+ t1$NumericSurvived <- ifelse(t1$Survived == "No", 0, 1)
+ t2 <- t1[-4]
+ df <- suppressWarnings(createDataFrame(t2))
+ m <- spark.naiveBayes(df, NumericSurvived ~ ., smoothing = 0.0)
+ s <- summary(m)
+ expect_equal(as.double(s$apriori[1, 1]), 0.5833333, tolerance = 1e-6)
+ expect_equal(sum(s$apriori), 1)
+ expect_equal(as.double(s$tables[1, "Age_Adult"]), 0.5714286, tolerance = 1e-6)
})
test_that("spark.survreg", {