first pass at spark execution support

topepo · topepo · commit b5cf7d744cb2 · 2018-04-04T13:47:40.000-04:00
diff --git a/R/fit.R b/R/fit.R
@@ -151,7 +151,7 @@ fit_formula <- function(object, formula, data, engine = engine, control, ...) {
         env = current_env()
       )
   } else {
-    if(object$method$interface %in% c("data.frame", "matrix")) {
+    if(object$method$interface %in% c("data.frame", "matrix", "spark")) {
       res <- formula_to_xy(object = object, formula = formula, data = data, control)
     } else {
       stop("I don't know about the ",
@@ -162,31 +162,54 @@ fit_formula <- function(object, formula, data, engine = engine, control, ...) {
   res
 }
 
+
+xy_to_xy <- function(object, x, y, control, ...) {
+  fit_expr <- object$method$fit_call
+  fit_expr[["x"]] <- quote(x)
+  fit_expr[["y"]] <- quote(y)
+    eval_mod(
+      fit_expr,
+      capture = control$verbosity == 0,
+      catch = control$catch,
+      env = current_env()
+    )
+}
+xy_to_matrix <- function(object, x, y, control, ...) {
+  if (object$method$interface == "matrix" && !is.matrix(x))
+    x <- as.matrix(x)
+  xy_to_xy(object, x, y, control, ...)
+}
+xy_to_df <- function(object, x, y, control, ...) {
+  if (object$method$interface == "data.frame" && !is.data.frame(x))
+    x <- as.data.frame(x)
+  xy_to_xy(object, x, y, control, ...)
+}
+xy_to_spark <- function(object, x, y, control, ...) {
+  sdf <- sparklyr::sdf_bind_cols(x, y)
+  fit_expr <- object$method$fit_call
+  fit_expr[["x"]] <- quote(sdf)
+  fit_expr[["features_col"]] <- quote(colnames(x))
+  fit_expr[["label_col"]] <- quote(colnames(y))
+    eval_mod(
+      fit_expr,
+      capture = control$verbosity == 0,
+      catch = control$catch,
+      env = current_env()
+    )
+}
+
+
 fit_xy <- function(object, x, y, control, ...) {
   opts <- quos(...)
 
-  # Look up the model's interface (e.g. formula, recipes, etc)
-  # and delegate to the connector functions (`xy_to_formula` etc)
-  if(object$method$interface == "formula") {
-    res <- xy_to_formula(object = object, x = x, y = y, control)
-  } else {
-    if(object$method$interface %in% c("data.frame", "matrix")) {
-      fit_expr <- object$method$fit_call
-      fit_expr[["x"]] <- quote(x)
-      fit_expr[["y"]] <- quote(y)
-      res <-
-        eval_mod(
-          fit_expr,
-          capture = control$verbosity == 0,
-          catch = control$catch,
-          env = current_env()
-        )
-    } else {
-      stop("I don't know about the ",
-           object$method$interface, " interface.",
-           call. = FALSE)
-    }
-  }
+  res <- switch(
+    object$method$interface,
+    formula = xy_to_formula(object = object, x = x, y = y, control, ...),
+    matrix = xy_to_matrix(object = object, x = x, y = y, control, ...),
+    data.frame = xy_to_df(object = object, x = x, y = y, control, ...),
+    spark = xy_to_spark(object = object, x = x, y = y, control, ...),
+    stop("Unknown interface")
+  )
   res
 }
 
@@ -394,11 +417,11 @@ has_both_or_none <- function(a, b)
 check_interface <- function(formula, recipe, x, y, data, cl) {
   inher(formula, "formula", cl)
   inher(recipe, "recipe", cl)
-  inher(x, c("data.frame", "matrix"), cl)
+  inher(x, c("data.frame", "matrix", "tbl_spark"), cl)
   # `y` can be a vector (which is not a class), or a factor (which is not a vector)
   if(!is.null(y) && !is.vector(y))
-    inher(y, c("data.frame", "matrix", "factor"), cl)
-  inher(data, c("data.frame", "matrix"), cl)
+    inher(y, c("data.frame", "matrix", "factor", "tbl_spark"), cl)
+  inher(data, c("data.frame", "matrix", "tbl_spark"), cl)
 
   x_interface <- !is.null(x) & !is.null(y)
   rec_interface <- !is.null(recipe) & !is.null(data)
diff --git a/R/logistic_reg_constr.R b/R/logistic_reg_constr.R
@@ -14,14 +14,14 @@ logistic_reg_engines <- data.frame(
   glm    = TRUE,
   glmnet = TRUE,
   spark  = TRUE,
-  stan   = TRUE,  
+  stan   = TRUE,
   row.names =  c("classification")
 )
 
 ###################################################################
 
 #' @importFrom stats binomial
-logistic_reg_glm_constr <- 
+logistic_reg_glm_constr <-
   function(
     formula = missing_arg(),
     family = binomial,
@@ -42,10 +42,10 @@ logistic_reg_glm_constr <-
   ) {
     libs <- "stats"
     interface <- "formula"
-    protect = c("formula", "data", "weights")   
+    protect = c("formula", "data", "weights")
     has_dots <- TRUE
     fit_name <- "glm"
-    fit_args <- 
+    fit_args <-
       enexprs(
         formula = formula,
         family = family,
@@ -64,7 +64,7 @@ logistic_reg_glm_constr <-
         y = y,
         contrasts = contrasts
       )
-    res <- 
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -78,7 +78,7 @@ logistic_reg_glm_constr <-
     res
   }
 
-logistic_reg_glmnet_constr <- 
+logistic_reg_glmnet_constr <-
   function(
     x = as.matrix(x),
     y = missing_arg(),
@@ -105,11 +105,11 @@ logistic_reg_glmnet_constr <-
     type.multinomial = c("ungrouped", "grouped")
   ) {
     libs <- "glmnet"
-    interface <- "data.frame"
-    protect = c("x", "y", "weights", "family")   
+    interface <- "matrix"
+    protect = c("x", "y", "weights", "family")
     has_dots <- FALSE
     fit_name <- "glmnet"
-    fit_args <- 
+    fit_args <-
       enexprs(
         x = x,
         y = y,
@@ -135,7 +135,7 @@ logistic_reg_glmnet_constr <-
         standardize.response = standardize.response,
         type.multinomial = type.multinomial
       )
-    res <- 
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -149,7 +149,7 @@ logistic_reg_glmnet_constr <-
     res
   }
 
-logistic_reg_stan_constr <- 
+logistic_reg_stan_constr <-
   function(
     formula = missing_arg(),
     family = binomial(),
@@ -173,10 +173,10 @@ logistic_reg_stan_constr <-
 ) {
     libs <- "rstanarm"
     interface <- "formula"
-    protect = c("formula", "data", "weights")   
+    protect = c("formula", "data", "weights")
     has_dots <- TRUE
     fit_name <- "stan_glm"
-    fit_args <- 
+    fit_args <-
       enexprs(
         formula = formula,
         family = family,
@@ -198,7 +198,7 @@ logistic_reg_stan_constr <-
         QR = QR,
         sparse = sparse
       )
-    res <- 
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -213,7 +213,7 @@ logistic_reg_stan_constr <-
   }
 
 
-logistic_reg_spark_constr <- 
+logistic_reg_spark_constr <-
   function(
     x = missing_arg(),
     formula = NULL,
@@ -239,11 +239,11 @@ logistic_reg_spark_constr <-
     uid = random_string("logistic_regression_")
   ) {
   libs <- "sparklyr"
-  interface <- "formula"
-  protect = c("features_col", "label_col", "x", "weight_col")   
+  interface <- "spark"
+  protect = c("features_col", "label_col", "x", "weight_col")
   has_dots <- TRUE
   fit_name <- "ml_logistic_regression"
-  fit_args <- 
+  fit_args <-
     enexprs(
       x = x,
       formula = formula,
@@ -268,7 +268,7 @@ logistic_reg_spark_constr <-
       raw_prediction_col = raw_prediction_col,
       uid = uid
     )
-  res <- 
+  res <-
     list(
       library = libs,
       interface = interface,
diff --git a/R/rand_forest_constr.R b/R/rand_forest_constr.R
@@ -2,7 +2,7 @@
 rand_forest_arg_key <- data.frame(
   randomForest = c("mtry", "ntree", "nodesize"),
   ranger = c("mtry", "num.trees", "min.node.size"),
-  spark = 
+  spark =
     c("feature_subset_strategy", "num_trees", "min_instances_per_node"),
   stringsAsFactors = FALSE,
   row.names =  c("mtry", "trees", "min_n")
@@ -19,7 +19,7 @@ rand_forest_engines <- data.frame(
 
 ###################################################################
 
-rand_forest_ranger_constr <- 
+rand_forest_ranger_constr <-
   function(
     formula = missing_arg(),
     data = missing_arg(),
@@ -52,10 +52,10 @@ rand_forest_ranger_constr <-
   ) {
     libs <- "ranger"
     interface <- "formula"
-    protect = c("formula", "data", "case.weights")   
+    protect = c("formula", "data", "case.weights")
     has_dots <- FALSE
     fit_name <- "ranger"
-    fit_args <- 
+    fit_args <-
       enexprs(
         formula = formula,
         data = data,
@@ -86,7 +86,7 @@ rand_forest_ranger_constr <-
         status.variable.name = status.variable.name,
         classification = classification
       )
-    res <- 
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -100,13 +100,13 @@ rand_forest_ranger_constr <-
     res
   }
 
-rand_forest_randomForest_constr <- 
+rand_forest_randomForest_constr <-
   function(
-    x  = missing_arg(), 
-    y = missing_arg(),  
-    xtest = NULL, 
-    ytest = NULL, 
-    ntree = 500, 
+    x  = missing_arg(),
+    y = missing_arg(),
+    xtest = NULL,
+    ytest = NULL,
+    ntree = 500,
     mtry = if (!is.null(y) && !is.factor(y))
       max(floor(ncol(x) / 3), 1)
     else
@@ -134,7 +134,7 @@ rand_forest_randomForest_constr <-
     keep.forest = !is.null(y) && is.null(xtest),
     corr.bias = FALSE,
     keep.inbag = FALSE
-  ) 
+  )
   {
     libs <- "randomForest"
     interface <- "data.frame"
@@ -166,8 +166,8 @@ rand_forest_randomForest_constr <-
         keep.forest = keep.forest,
         corr.bias = corr.bias,
         keep.inbag = keep.inbag
-      ) 
-    res <- 
+      )
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -178,11 +178,11 @@ rand_forest_randomForest_constr <-
         fit_call = NULL
       )
     class(res) <- c("rand_forest_constr")
-    res  
+    res
   }
 
 
-rand_forest_spark_constr <- 
+rand_forest_spark_constr <-
   function(
     x = missing_arg(),
     formula = NULL,
@@ -211,7 +211,7 @@ rand_forest_spark_constr <-
   )
   {
     libs <- "sparklyr"
-    interface <- "tbl_spark"
+    interface <- "spark"
     protect = c("x", "features_col", "label_col", "type")
     has_dots <- TRUE
     fit_name <- "ml_random_forest"
@@ -241,8 +241,8 @@ rand_forest_spark_constr <-
         uid = uid,
         response = response,
         features = features
-      ) 
-    res <- 
+      )
+    res <-
       list(
         library = libs,
         interface = interface,
@@ -253,7 +253,7 @@ rand_forest_spark_constr <-
         fit_call = NULL
       )
     class(res) <- c("rand_forest_constr")
-    res  
+    res
   }
 
 ###################################################################
diff --git a/tests/testthat/test_logistic_reg.R b/tests/testthat/test_logistic_reg.R
@@ -315,17 +315,16 @@ test_that('glmnet execution', {
   #   regexp = NA
   # )
 
-  # fails during R CMD check but works outside of that
-  # expect_error(
-  #   fit(
-  #     lc_basic,
-  #     engine = "glmnet",
-  #     control = ctrl,
-  #     x = lending_club[, num_pred],
-  #     y = lending_club$Class
-  #   ),
-  #   regexp = NA
-  # )
+  expect_error(
+    fit(
+      lc_basic,
+      engine = "glmnet",
+      control = ctrl,
+      x = lending_club[, num_pred],
+      y = lending_club$Class
+    ),
+    regexp = NA
+  )
 
   # fails because `glment` requires a matrix
   # expect_error(
