fixed spark models so that class labels (not integers) are in colnames

Author: topepo

R/aaa_spark_helpers.R

@@ -3,12 +3,10 @@
 #' @importFrom dplyr starts_with rename rename_at vars funs
 format_spark_probs <- function(results, object) {
   results <- dplyr::select(results, starts_with("probability_"))
-  results <- dplyr::rename_at(
-    results,
-    vars(starts_with("probability_")),
-    funs(gsub("probability", "pred", .))
-  )
-  results
+  p <- ncol(results)
+  lvl <- paste0("probability_", 0:(p - 1))
+  names(lvl) <- paste0("pred_", object$fit$.index_labels)
+  results %>% rename(!!!syms(lvl))
 }
 
 format_spark_class <- function(results, object) {

tests/testthat/test_logistic_reg_spark.R

@@ -79,7 +79,7 @@ test_that('spark execution', {
     regexp = NA
   )
 
-  expect_equal(colnames(spark_class_prob), c("pred_0", "pred_1"))
+  expect_equal(colnames(spark_class_prob), c("pred_Yes", "pred_No"))
 
   expect_equivalent(
     as.data.frame(spark_class_prob),

tests/testthat/test_multinom_reg_spark.R

@@ -69,7 +69,7 @@ test_that('spark execution', {
 
   expect_equal(
     colnames(spark_class_prob),
-    c("pred_0", "pred_1", "pred_2")
+    c("pred_versicolor", "pred_virginica", "pred_setosa")
   )
 
   expect_equivalent(