classification unit test

Author: vectorijk

R/pkg/R/mllib.R

@@ -1474,7 +1474,7 @@ print.summary.KSTest <- function(x, ...) {
 #' df <- createDataFrame(longley)
 #'
 #' # fit a Decision Tree Regression Model
-#' model <- spark.decisionTree(data, Employed~., "regression", maxDepth=5, maxBins=16)
+#' model <- spark.decisionTree(data, Employed~., type = "regression", maxDepth = 5, maxBins = 16)
 #'
 #' # get the summary of the model
 #' summary(model)
@@ -1579,7 +1579,7 @@ setMethod("summary", signature(object = "DecisionTreeRegressionModel"),
             features <- callJMethod(jobj, "features")
             depth <- callJMethod(jobj, "depth")
             numNodes <- callJMethod(jobj, "numNodes")
-            ans <- list(features = features, depth = depth, numNodes = numNodes)
+            ans <- list(features = features, depth = depth, numNodes = numNodes, jobj = jobj)
             class(ans) <- "summary.DecisionTreeRegressionModel"
             ans
           })
@@ -1594,15 +1594,17 @@ setMethod("summary", signature(object = "DecisionTreeRegressionModel"),
 #' @export
 #' @note summary(DecisionTreeRegressionModel) since 2.1.0
 setMethod("summary", signature(object = "DecisionTreeClassificationModel"),
-function(object, ...) {
-    jobj <- object@jobj
-    features <- callJMethod(jobj, "features")
-    depth <- callJMethod(jobj, "depth")
-    numNodes <- callJMethod(jobj, "numNodes")
-    ans <- list(features = features, depth = depth, numNodes = numNodes)
-    class(ans) <- "summary.DecisionTreeClassificationModel"
-    ans
-})
+          function(object, ...) {
+              jobj <- object@jobj
+              features <- callJMethod(jobj, "features")
+              depth <- callJMethod(jobj, "depth")
+              numNodes <- callJMethod(jobj, "numNodes")
+              numClasses <- callJMethod(jobj, "numClasses")
+              ans <- list(features = features, depth = depth,
+                          numNodes = numNodes, numClasses = numClasses, jobj = jobj)
+              class(ans) <- "summary.DecisionTreeClassificationModel"
+              ans
+          })
 
 #  Prints the summary of Decision Tree Regression Model
 
@@ -1611,11 +1613,11 @@ function(object, ...) {
 #' @export
 #' @note print.summary.DecisionTreeRegressionModel since 2.1.0
 print.summary.DecisionTreeRegressionModel <- function(x, ...) {
-  jobj <- x@jobj
-  summaryStr <- callJMethod(jobj, "summary")
-  cat(summaryStr, "\n")
-  invisible(x)
-  }
+    jobj <- x$jobj
+    summaryStr <- callJMethod(jobj, "summary")
+    cat(summaryStr, "\n")
+    invisible(x)
+}
 
 #  Prints the summary of Decision Tree Classification Model
 
@@ -1624,7 +1626,7 @@ print.summary.DecisionTreeRegressionModel <- function(x, ...) {
 #' @export
 #' @note print.summary.DecisionTreeClassificationModel since 2.1.0
 print.summary.DecisionTreeClassificationModel <- function(x, ...) {
-    jobj <- x@jobj
+    jobj <- x$jobj
     summaryStr <- callJMethod(jobj, "summary")
     cat(summaryStr, "\n")
     invisible(x)

R/pkg/inst/tests/testthat/test_mllib.R

@@ -821,4 +821,29 @@ test_that("spark.decisionTree Regression", {
   unlink(modelPath)
 })
 
+test_that("spark.decisionTree Classification", {
+  data <- suppressWarnings(createDataFrame(iris))
+  model <- spark.decisionTree(data, Species ~ Petal_Length + Petal_Width, "classification",
+                              maxDepth = 5, maxBins = 16)
+
+  #Test summary
+  stats <- summary(model)
+  expect_equal(stats$depth, 5)
+  expect_equal(stats$numNodes, 19)
+  expect_equal(stats$numClasses, 3)
+
+  # Test model save/load
+  modelPath <- tempfile(pattern = "spark-decisionTreeClassification", fileext = ".tmp")
+  write.ml(model, modelPath)
+  expect_error(write.ml(model, modelPath))
+  write.ml(model, modelPath, overwrite = TRUE)
+  model2 <- read.ml(modelPath)
+  stats2 <- summary(model2)
+  expect_equal(stats$depth, stats2$depth)
+  expect_equal(stats$numNodes, stats2$numNodes)
+  expect_equal(stats$numClasses, stats2$numClasses)
+
+  unlink(modelPath)
+})
+
 sparkR.session.stop()

mllib/src/main/scala/org/apache/spark/ml/r/DecisionTreeClassifierWrapper.scala

@@ -38,6 +38,12 @@ private[r] class DecisionTreeClassifierWrapper private (
   private val DTModel: DecisionTreeClassificationModel =
     pipeline.stages(1).asInstanceOf[DecisionTreeClassificationModel]
 
+  lazy val depth: Int = DTModel.depth
+  lazy val numNodes: Int = DTModel.numNodes
+  lazy val numClasses: Int = DTModel.numClasses
+
+  def summary: String = DTModel.toDebugString
+
   def transform(dataset: Dataset[_]): DataFrame = {
     pipeline.transform(dataset)
       .drop(DTModel.getFeaturesCol)