vctrs based slice()

R/tbl-df.r

@@ -221,11 +221,79 @@ slice.tbl_df <- function(.data, ..., .preserve = FALSE) {
     return(.data)
   }
 
+  rows <- group_rows(.data)
+  mask <- groupwise_data_mask(.data, rows)
+  caller <- caller_env()
+
+  old_group_size <- context_env[["..group_size"]]
+  old_group_number <- context_env[["..group_number"]]
+  on.exit({
+    context_env[["..group_size"]] <- old_group_size
+    context_env[["..group_number"]] <- old_group_number
+  })
+
   quo <- quo(c(!!!dots))
-  out <- slice_impl(.data, quo)
-  if (!.preserve && is_grouped_df(.data)) {
-    out <- regroup(out)
+
+  slice_indices <- new_list(length(rows))
+  new_rows <- new_list(length(rows))
+  k <- 1L
+
+  for (group in seq_along(rows)) {
+    current_rows <- rows[[group]]
+    mask$.set_current_group(group)
+    n <- length(current_rows)
+    if (n == 0L) {
+      new_rows[[group]] <- integer()
+      next
+    }
+    context_env[["..group_size"]] <- n
+    context_env[["..group_number"]] <- group
+
+    res <- eval_tidy(quo, mask, env = caller)
+
+    if (is.logical(res) && all(is.na(res))) {
+      res <- integer()
+    } else if (is.numeric(res)) {
+      res <- vec_cast(res, integer())
+    } else if (!is.integer(res)) {
+      abort(
+        "slice() expressions should return indices (positive or negative integers)",
+        "dplyr_slice_incompatible"
+      )
+    }
+
+    if (length(res) == 0L) {
+      # nothing to do
+    } else if(all(res >= 0, na.rm = TRUE)) {
+      res <- res[!is.na(res) & res <= length(current_rows) & res > 0]
+    } else if (all(res <= 0, na.rm = TRUE)) {
+      res <- setdiff(seq_along(current_rows), -res)
+    } else {
+      abort(
+        "slice() expressions should return either all positive or all negative",
+        "dplyr_slice_ambiguous"
+      )
+    }
+
+    slice_indices[[group]] <- current_rows[res]
+    new_k <- k + length(res)
+    new_rows[[group]] <- seq2(k, new_k - 1L)
+    k <- new_k
   }
+  all_slice_indices <- vec_c(!!!slice_indices, .ptype = integer())
+
+  out <- vec_slice(.data, all_slice_indices)
+
+  if (is_grouped_df(.data)) {
+    new_groups <- group_data(.data)
+    new_groups$.rows <- new_list_of(new_rows, ptype = integer())
+    attr(out, "groups") <- new_groups
+
+    if (!.preserve) {
+      out <- regroup(out)
+    }
+  }
+
   out
 }
 #' @export

tests/testthat/test-slice.r

@@ -22,22 +22,17 @@ test_that("slice silently ignores out of range values (#226)", {
 test_that("slice works with negative indices", {
   res <- slice(mtcars, -(1:2))
   exp <- tail(mtcars, -2)
-  expect_equal(names(res), names(exp))
-  for (col in names(res)) {
-    expect_equal(res[[col]], exp[[col]])
-  }
+  expect_equivalent(res, exp)
 })
 
 test_that("slice forbids positive and negative together", {
   expect_error(
     mtcars %>% slice(c(-1, 2)),
-    "Found 1 positive indices and 1 negative indices",
-    fixed = TRUE
+    class = "dplyr_slice_ambiguous"
   )
   expect_error(
     mtcars %>% slice(c(2:3, -1)),
-    "Found 2 positive indices and 1 negative indices",
-    fixed = TRUE
+    class = "dplyr_slice_ambiguous"
   )
 })
 
@@ -86,8 +81,14 @@ test_that("slice handles NA (#1235)", {
 test_that("slice handles logical NA (#3970)", {
   df <- tibble(x = 1:3)
   expect_equal(nrow(slice(df, NA)), 0L)
-  expect_error(slice(df, TRUE))
-  expect_error(slice(df, FALSE))
+  expect_error(
+    slice(df, TRUE),
+    class = "dplyr_slice_incompatible"
+  )
+  expect_error(
+    slice(df, FALSE),
+    class = "dplyr_slice_incompatible"
+  )
 })
 
 test_that("slice handles empty data frames (#1219)", {
@@ -116,12 +117,6 @@ test_that("slice works with zero-column data frames (#2490)", {
   )
 })
 
-test_that("slice works under gctorture2", {
-  x <- tibble(y = 1:10)
-  with_gctorture2(999, x2 <- slice(x, 1:10))
-  expect_identical(x, x2)
-})
-
 test_that("slice correctly computes positive indices from negative indices (#3073)", {
   x <- tibble(y = 1:10)
   expect_identical(slice(x, -10:-30), tibble(y = 1:9))
@@ -159,12 +154,6 @@ test_that("slice skips 0 (#3313)", {
   expect_identical(slice(d, c(0, -1)), slice(d, -1))
 })
 
-test_that("slice is not confused about dense groups (#3753)",{
-  df <- tibble(row = 1:3)
-  expect_equal(slice(df, c(2,1,3))$row, c(2L,1L,3L))
-  expect_equal(slice(df, c(1,1,1))$row, rep(1L, 3))
-})
-
 test_that("slice accepts ... (#3804)", {
   expect_equal(slice(mtcars, 1, 2), slice(mtcars, 1:2))
   expect_equal(slice(mtcars, 1, n()), slice(mtcars, c(1, nrow(mtcars))))
@@ -187,8 +176,3 @@ test_that("slice does not evaluate the expression in empty groups (#1438)", {
   )
   expect_equal(nrow(res), 3L)
 })
-
-test_that("column_subset() falls back to R indexing on esoteric data types (#4128)", {
-  res <- slice(tibble::enframe(formals(rnorm)), 2:3)
-  expect_identical(res, tibble(name = c("mean", "sd"), value = list(0, 1)))
-})