Merge pull request #396 from tidyverse/rank-funs

Translate `min_rank`, `dense_rank`, `percent_rank`, & `cume_dist`
tidyverse · Nov 1, 2022 · 5d71f4c · 5d71f4c
2 parents b42d698 + e8ca0c9
commit 5d71f4c
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 5 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,8 @@
 # dtplyr (development version)
 
+* `min_rank()`, `dense_rank()`, `percent_rank()`, & `cume_dist()` are now translated
+  to their `data.table` equivalents (#396)
+
 * `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 

diff --git a/R/tidyeval.R b/R/tidyeval.R
@@ -152,9 +152,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
   } else if (is_call(x, "cur_group_rows")) {
     quote(.I)
   } else if (is_call(x, "desc")) {
-    if (!has_length(x, 2L)) {
-      abort("`desc()` expects exactly one argument.")
-    }
+    check_one_arg(x)
     x[[1]] <- sym("-")
     x[[2]] <- dt_squash(x[[2]], env, data, j)
     x
@@ -212,6 +210,24 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
   } else if (is_call(x, "row_number", n = 1)) {
     arg <- dt_squash(x[[2]], env, data, j = j)
     expr(frank(!!arg, ties.method = "first", na.last = "keep"))
+  } else if (is_call(x, "min_rank")) {
+    check_one_arg(x)
+    arg <- dt_squash(x[[2]], env, data, j = j)
+    expr(frank(!!arg, ties.method = "min", na.last = "keep"))
+  } else if (is_call(x, "dense_rank")) {
+    check_one_arg(x)
+    arg <- dt_squash(x[[2]], env, data, j = j)
+    expr(frank(!!arg, ties.method = "dense", na.last = "keep"))
+  } else if (is_call(x, "percent_rank")) {
+    check_one_arg(x)
+    arg <- dt_squash(x[[2]], env, data, j = j)
+    frank_expr <- expr((frank(!!arg, ties.method = "min", na.last = "keep") - 1))
+    expr(!!frank_expr / (sum(!is.na(!!arg)) - 1))
+  } else if (is_call(x, "cume_dist")) {
+    check_one_arg(x)
+    arg <- dt_squash(x[[2]], env, data, j = j)
+    frank_expr <- expr(frank(!!arg, ties.method = "max", na.last = "keep"))
+    expr(!!frank_expr / sum(!is.na(!!arg)))
   } else if (is.function(x[[1]]) || is_call(x, "function")) {
     simplify_function_call(x, env, data, j = j)
   } else if (is_call(x, c("glue", "str_glue")) && j) {
@@ -301,3 +317,10 @@ fun_name <- function(fun) {
 
   NULL
 }
+
+check_one_arg <- function(x) {
+  fun <- as_name(x[[1]])
+  if (!has_length(x, 2L)) {
+    abort(glue("`{fun}()` expects exactly one argument."))
+  }
+}
diff --git a/tests/testthat/_snaps/tidyeval.md b/tests/testthat/_snaps/tidyeval.md
@@ -15,6 +15,6 @@
     Code
       capture_dot(df, desc(a, b))
     Condition
-      Error in `dt_squash_call()`:
+      Error in `check_one_arg()`:
       ! `desc()` expects exactly one argument.
 
diff --git a/tests/testthat/test-tidyeval.R b/tests/testthat/test-tidyeval.R
@@ -265,7 +265,7 @@ test_that("n() is equivalent to .N", {
   )
 })
 
-test_that("row_number() is equivalent .I", {
+test_that("row_number() is equivalent seq_len(.N)", {
   dt <- lazy_dt(data.frame(g = c(1, 1, 2), x = 1:3))
 
   expect_equal(
@@ -286,6 +286,21 @@ test_that("row_number(x) is equivalent to rank", {
   )
 })
 
+test_that("ranking functions are translated", {
+  df <- lazy_dt(tibble(x = c(1, 2, NA, 1, 0, NaN)))
+
+  res <- df %>%
+    mutate(percent_rank = percent_rank(x),
+           min_rank = min_rank(x),
+           dense_rank = dense_rank(x),
+           cume_dist = cume_dist(x))
+
+  expect_equal(pull(res, percent_rank), c(1 / 3, 1, NA, 1 / 3, 0, NA))
+  expect_equal(pull(res, min_rank), c(2L, 4L, NA, 2L, 1L, NA))
+  expect_equal(pull(res, dense_rank), c(2L, 3L, NA, 2L, 1L, NA))
+  expect_equal(pull(res, cume_dist), c(.75, 1, NA, .75, .25, NA))
+})
+
 test_that("scoped verbs produce nice output", {
   dt <- lazy_dt(data.table(x = 1:5), "DT")