A few more tests, slightly more guard rails for unimplemented types

apache · Jul 15, 2021 · ee206d5 · ee206d5
1 parent a823c34
commit ee206d5
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 18 deletions.
diff --git a/r/R/dplyr-functions.R b/r/R/dplyr-functions.R
@@ -568,8 +568,20 @@ nse_funcs$if_else <- function(condition, true, false, missing = NULL){
   # ARROW-13186 might make this easier with a more robust way.
   # TODO: do this ^^^
 
-  if (inherits(true, "character") || inherits(false, "character")) {
-    stop("`true` and `false` character values not yet supported in Arrow")
+  # if_else only supports boolean, numeric, or temporal types right now
+  # TODO: remove when ARROW-12955 merges
+  # If true/false are R types, we can use `is.*` directly
+  invalid_r_types <- is.character(true) || is.character(false) || is.list(true) ||
+    is.list(false) || is.factor(true) || is.factor(false)
+  # However, if they are expressions, we need to use the functions from nse_funcs
+  invalid_expression_types_true <- inherits(true, "Expression") && (
+    nse_funcs$is.character(true) || nse_funcs$is.list(true) || nse_funcs$is.factor(true)
+  )
+  invalid_expression_types_false <- inherits(false, "Expression") && (
+    nse_funcs$is.character(false) || nse_funcs$is.list(false) || nse_funcs$is.factor(false)
+  )
+  if (invalid_r_types | invalid_expression_types_true | invalid_expression_types_false) {
+    stop("`true` and `false` character values not yet supported in Arrow", call. = FALSE)
   }
 
   build_expr("if_else", condition, true, false)

diff --git a/r/tests/testthat/test-dplyr.R b/r/tests/testthat/test-dplyr.R
@@ -947,61 +947,96 @@ test_that("abs()", {
 })
 
 test_that("if_else and ifelse", {
-  df <- tibble(x = c(-127, -10, -1, -0 , 0, 1, 10, 127, NA))
-
   expect_dplyr_equal(
     input %>%
       mutate(
-        y = if_else(x > 0, 1, 0)
+        y = if_else(int > 5, 1, 0)
       ) %>% collect(),
-    df
+    example_data
   )
 
   expect_dplyr_equal(
     input %>%
       mutate(
-        y = if_else(x > 0, x, 0)
+        y = if_else(int > 5, int, 0L)
       ) %>% collect(),
-    df
+    example_data
   )
 
   expect_error(
-    Table$create(df) %>%
+    Table$create(example_data) %>%
       mutate(
-        y = if_else(x > 0, 1, FALSE)
+        y = if_else(int > 5, 1, FALSE)
       ) %>% collect(),
     'NotImplemented: Function if_else has no kernel matching input types'
   )
 
   expect_dplyr_equal(
     input %>%
       mutate(
-        y = if_else(x > 0, 1, NA_real_)
+        y = if_else(int > 5, 1, NA_real_)
       ) %>% collect(),
-    df
+    example_data
   )
 
   expect_dplyr_equal(
     input %>%
       mutate(
-        y = ifelse(x > 0, 1, 0)
+        y = ifelse(int > 5, 1, 0)
       ) %>% collect(),
-    df
+    example_data
   )
 
   expect_dplyr_equal(
     input %>%
       mutate(
-        y = ifelse(x > 0, x, 0)
+        y = ifelse(dbl > 5, TRUE, FALSE)
       ) %>% collect(),
-    df
+    example_data
+  )
+
+  expect_dplyr_equal(
+    input %>%
+      mutate(
+        y = ifelse(chr %in% letters[1:3], 1L, 3L)
+      ) %>% collect(),
+    example_data
+  )
+
+  # TODO: this should not warn / pull into R, once ARROW-12955 merges
+  expect_dplyr_equal(
+    input %>%
+      mutate(
+        y = if_else(int > 5, "one", "zero")
+      ) %>% collect(),
+    example_data,
+    warn = TRUE
+  )
+
+  # TODO: this should not warn / pull into R, once ARROW-12955 merges
+  expect_dplyr_equal(
+    input %>%
+      mutate(
+        y = if_else(int > 5, chr, chr)
+      ) %>% collect(),
+    example_data,
+    warn = TRUE
+  )
+
+  expect_dplyr_equal(
+    input %>%
+      mutate(
+        y = if_else(int > 5, fct, factor("a"))
+      ) %>% collect(),
+    example_data,
+    warn = TRUE
   )
 
   skip("TODO: could? should? we support the autocasting in ifelse")
   expect_dplyr_equal(
     input %>%
-      mutate(y = ifelse(x > 0, 1, FALSE)) %>%
+      mutate(y = ifelse(int > 5, 1, FALSE)) %>%
       collect(),
-    df
+    example_data
   )
 })