add format helper

Author: hfrick

R/predict.R

@@ -222,37 +222,19 @@ check_pred_type <- function(object, type, ...) {
 #' @export
 
 format_num <- function(x) {
-  if (inherits(x, "tbl_spark"))
+  if (inherits(x, "tbl_spark")) {
     return(x)
-
-  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
-    x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred", names(x)))) {
-      names(x) <- paste0(".pred_", names(x))
-    }
-  } else {
-    x <- tibble(.pred = unname(x))
   }
-
-  x
+  ensure_parsnip_format(x, ".pred", overwrite = FALSE)
 }
 
 #' @rdname format-internals
 #' @export
 format_class <- function(x) {
-  if (inherits(x, "tbl_spark"))
+  if (inherits(x, "tbl_spark")) {
     return(x)
-
-  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
-    x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred_class", names(x)))) {
-      names(x) <- ".pred_class"
-    }
-  } else {
-    x <- tibble(.pred_class = unname(x))
   }
-
-  x
+  ensure_parsnip_format(x, ".pred_class")
 }
 
 #' @rdname format-internals
@@ -269,64 +251,45 @@ format_classprobs <- function(x) {
 #' @rdname format-internals
 #' @export
 format_time <- function(x) {
-  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
-    x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred_time", names(x)))) {
-      names(x) <- paste0(".pred_time_", names(x))
-    }
-  } else {
-    x <- tibble(.pred_time = unname(x))
-  }
-
-  x
+  ensure_parsnip_format(x, ".pred_time", overwrite = FALSE)
 }
 
 #' @rdname format-internals
 #' @export
 format_survival <- function(x) {
-  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
-    x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred", names(x)))) {
-      names(x) <- ".pred"
-    }
-  } else {
-    x <- tibble(.pred = unname(x))
-  }
-
-  x
+  ensure_parsnip_format(x, ".pred")
 }
 
 #' @rdname format-internals
 #' @export
 format_linear_pred <- function(x) {
-  if (inherits(x, "tbl_spark"))
+  if (inherits(x, "tbl_spark")){
     return(x)
-
-  if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
-    x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred_linear_pred", names(x)))) {
-      names(x) <- ".pred_linear_pred"
-    }
-  } else {
-    x <- tibble(.pred_linear_pred = unname(x))
   }
-
-  x
+  ensure_parsnip_format(x, ".pred_linear_pred")
 }
 
 #' @rdname format-internals
 #' @export
 format_hazard <- function(x) {
+  ensure_parsnip_format(x, ".pred")
+}
+
+ensure_parsnip_format <- function(x, col_name, overwrite = TRUE) {
   if (isTRUE(ncol(x) > 1) | is.data.frame(x)) {
     x <- as_tibble(x, .name_repair = "minimal")
-    if (!any(grepl("^\\.pred", names(x)))) {
-      names(x) <- ".pred"
+    if (!any(grepl(paste0("^\\", col_name), names(x)))) {
+      if (overwrite) {
+        names(x) <- col_name
+      } else {
+        names(x) <- paste(col_name, names(x), sep = "_")
+      }
     }
+  } else {
+    x <- tibble(unname(x))
+    names(x) <- col_name
+    x
   }
-  else {
-    x <- tibble(.pred_hazard = unname(x))
-  }
-
   x
 }