Merge pull request #395 from tidyverse/step_setnames-NA

Correctly handle `NA` column names from `pivot_wider()`
tidyverse · Oct 21, 2022 · b42d698 · b42d698
2 parents 55e9298 + 7f3761e
commit b42d698
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 0 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,7 @@
 # dtplyr (development version)
 
+* `names_glue` now works in `pivot_wider()` when `names_from` contains `NA`s (#394)
+
 * `full_join()` now produces output with correctly named columns when a non-default 
   value for `suffix` is supplied. Previously the `suffix` argument was ignored (#382).
 

diff --git a/R/step-call-pivot_wider.R b/R/step-call-pivot_wider.R
@@ -74,6 +74,8 @@ pivot_wider.dtplyr_step <- function(data,
     new_vars <- as.character(new_vars)
   }
 
+  new_vars <- vctrs::vec_assign(new_vars, is.na(new_vars), "NA")
+
   if (!is.null(names_glue)) {
     glue_df <- as.data.table(distinct(ungroup(data), !!!syms(names_from)))
     glue_df <- vctrs::vec_rep(glue_df, length(values_from))

diff --git a/tests/testthat/test-step-call-pivot_longer.R b/tests/testthat/test-step-call-pivot_longer.R
@@ -151,3 +151,19 @@ test_that("informative errors on unsupported features", {
 
 })
 
+test_that("can pivot all cols to long", {
+  tbl <- tibble(x = 1:2, y = 3:4)
+  dt <- lazy_dt(tbl, "DT")
+  step <- pivot_longer(dt, x:y)
+  out <- collect(step)
+
+  expect_equal(
+    show_query(step),
+    expr(melt(DT, measure.vars = !!c("x", "y"), variable.name = "name",
+              variable.factor = FALSE))
+  )
+  expect_equal(step$vars, c("name", "value"))
+  expect_equal(out$name, c("x", "x", "y", "y"))
+  expect_equal(out$value, c(1, 2, 3, 4))
+})
+
diff --git a/tests/testthat/test-step-call-pivot_wider.R b/tests/testthat/test-step-call-pivot_wider.R
@@ -63,6 +63,33 @@ test_that("column with `...j` name can be used as `names_from`", {
   expect_equal(nrow(pv), 1)
 })
 
+test_that("correctly handles columns named NA, #394", {
+  df <- lazy_dt(tibble(x = c("a", "a"), y = c("a", NA), z = 1:2))
+
+  res <- df %>%
+    pivot_wider(names_from = y,
+                values_from = z,
+                names_glue = "{y}_new",
+                names_repair = "minimal") %>%
+    collect()
+  expect_named(res, c("x", "NA_new", "a_new"))
+
+  res <- df %>%
+    pivot_wider(names_from = y,
+                values_from = z,
+                names_glue = "{y}",
+                names_repair = "minimal") %>%
+    collect()
+  expect_named(res, c("x", "NA", "a"))
+
+  df <- lazy_dt(tibble(x = c('a', NA), y = 1:2))
+
+  res <- df %>%
+    pivot_wider(names_from = 'x', values_from = 'y') %>%
+    collect()
+  expect_named(res, c("NA", "a"))
+})
+
 # column names -------------------------------------------------------------
 
 test_that("names_glue affects output names", {