feat: Implement slice_head() (#640)

Author: krlmlr

R/head.R

@@ -9,7 +9,7 @@ head.duckplyr_df <- function(x, n = 6L, ...) {
     #' - with a negative `n`.
     #'
     #' These features fall back to [head()], see `vignette("fallback")` for details.
-    "Can't process negative n" = (n < 0),
+    "{.code slice_head(n = ...)} with negative values not supported" = (n < 0),
     {
       rel <- duckdb_rel_from_df(x)
       out_rel <- rel_limit(rel, n)

R/relational.R

@@ -27,7 +27,7 @@ rel_try <- function(call, rel, ...) {
 
   if (Sys.getenv("DUCKPLYR_FALLBACK_FORCE") == "TRUE") {
     stats$fallback <- stats$fallback + 1L
-    return()
+    return("Fallback enforced")
   }
 
   dots <- list(...)

R/slice_head-rd.R

@@ -0,0 +1,15 @@
+#' @title Subset rows using their positions
+#'
+#' @description  This is a method for the [dplyr::slice_head()] generic.
+#' `slice_head()` selects the first rows.
+#'
+#' @inheritParams dplyr::slice_head
+#' @examples
+#' library(duckplyr)
+#' df <- data.frame(x = 1:3)
+#' df <- slice_head(df, n = 2)
+#' df
+#' @seealso [dplyr::slice_head()]
+#' @rdname slice_head.duckplyr_df
+#' @name slice_head.duckplyr_df
+NULL

R/slice_head.R

@@ -1,10 +1,28 @@
 # Generated by 02-duckplyr_df-methods.R
+#' @rdname slice_head.duckplyr_df
 #' @export
 slice_head.duckplyr_df <- function(.data, ..., n, prop, by = NULL) {
+  if (!missing(n)) {
+    n_valid <- (n >= 0)
+  } else {
+    n_valid <- TRUE
+  }
+
   # Our implementation
   duckplyr_error <- rel_try(NULL,
-    "No relational implementation for {.code slice_head()}" = TRUE,
+    #' @section Fallbacks:
+    #' There is no DuckDB translation in `slice_head.duckplyr_df()`
+    #' - if `by` or `prop` is provided,
+    #' - with a negative `n`.
+    #'
+    #' These features fall back to [slice_head()], see `vignette("fallback")` for details.
+    "{.code slice_head(by = ...)} not supported" = !missing(by),
+    "{.code slice_head(prop = ...)} not supported" = !missing(prop),
+    "{.code slice_head(n = ...)} with negative values not supported" = !n_valid,
     {
+      rel <- duckdb_rel_from_df(.data)
+      out_rel <- rel_limit(rel, n)
+      out <- duckplyr_reconstruct(out_rel, .data)
       return(out)
     }
   )

R/tpch.R

@@ -77,7 +77,7 @@ tpch_02 <- function() {
       s_address, s_phone, s_comment
     ) %>%
     arrange(desc(s_acctbal), n_name, s_name, p_partkey) %>%
-    head(100)
+    slice_head(n = 100)
 
   res
 }
@@ -109,7 +109,7 @@ tpch_03 <- function() {
     summarise(revenue = sum(na.rm = TRUE, volume), .by = c(l_orderkey, o_orderdate, o_shippriority)) %>%
     select(l_orderkey, revenue, o_orderdate, o_shippriority) %>%
     arrange(desc(revenue), o_orderdate) %>%
-    head(10)
+    slice_head(n = 10)
   aggr
 }
 
@@ -450,7 +450,7 @@ tpch_10 <- function() {
       c_address, c_phone, c_comment
     ) %>%
     arrange(desc(revenue)) %>%
-    head(20)
+    slice_head(n = 20)
 
   res
 }

R/tpch2.R

@@ -186,7 +186,7 @@ tpch_18 <- function() {
       o_orderdate, o_totalprice, sum
     ) %>%
     arrange(desc(o_totalprice), o_orderdate) %>%
-    head(100)
+    slice_head(n = 100)
 }
 
 #' @autoglobal
@@ -305,7 +305,7 @@ tpch_21 <- function() {
     filter(n_name == "SAUDI ARABIA") %>%
     summarise(numwait = n(), .by = s_name) %>%
     arrange(desc(numwait), s_name) %>%
-    head(100)
+    slice_head(n = 100)
 }
 
 #' @autoglobal

_pkgdown.yml

@@ -49,6 +49,7 @@ reference:
   - arrange.duckplyr_df
   - distinct.duckplyr_df
   - filter.duckplyr_df
+  - slice_head.duckplyr_df
   - head.duckplyr_df
 
 - subtitle: Verbs that affect columns

man/slice_head.duckplyr_df.Rd

@@ -0,0 +1,33 @@
+diff --git b/R/slice_head.R a/R/slice_head.R
+--- b/R/slice_head.R
++++ a/R/slice_head.R
+@@ -1,10 +1,28 @@
+ # Generated by 02-duckplyr_df-methods.R
++#' @rdname slice_head.duckplyr_df
+ #' @export
+ slice_head.duckplyr_df <- function(.data, ..., n, prop, by = NULL) {
++  if (!missing(n)) {
++    n_valid <- (n >= 0)
++  } else {
++    n_valid <- TRUE
++  }
++
+   # Our implementation
+   duckplyr_error <- rel_try(NULL,
+-    "No relational implementation for {.code slice_head()}" = TRUE,
++    #' @section Fallbacks:
++    #' There is no DuckDB translation in `slice_head.duckplyr_df()`
++    #' - if `by` or `prop` is provided,
++    #' - with a negative `n`.
++    #'
++    #' These features fall back to [slice_head()], see `vignette("fallback")` for details.
++    "{.code slice_head(by = ...)} not supported" = !is.null(by),
++    "{.code slice_head(prop = ...)} not supported" = !missing(prop),
++    "{.code slice_head(n = ...)} with negative values not supported" = !n_valid,
+     {
++      rel <- duckdb_rel_from_df(.data)
++      out_rel <- rel_limit(rel, n)
++      out <- duckplyr_reconstruct(out_rel, .data)
+       return(out)
+     }
+   )

patch/slice_head.patch

@@ -2346,17 +2346,28 @@ test_that("as_duckplyr_df_impl() and slice()", {
   expect_identical(pre, post)
 })
 
-test_that("as_duckplyr_df_impl() and slice_head()", {
-  withr::local_envvar(DUCKPLYR_FORCE = "FALSE")
+test_that("as_duckplyr_df_impl() and slice_head(n = 2)", {
+  withr::local_envvar(DUCKPLYR_FALLBACK_FORCE = "TRUE")
+
+  # Data
+  test_df <- data.frame(a = 1:6 + 0, b = 2, g = rep(1:3, 1:3))
+
+  # Run
+  pre <- test_df %>% as_duckplyr_df_impl() %>% slice_head(n = 2)
+  post <- test_df %>% slice_head(n = 2) %>% as_duckplyr_df_impl()
+
+  # Compare
+  expect_identical(pre, post)
+})
 
-  skip("External vector?")
 
+test_that("as_duckplyr_df_impl() and slice_head(n = 2)", {
   # Data
   test_df <- data.frame(a = 1:6 + 0, b = 2, g = rep(1:3, 1:3))
 
   # Run
-  pre <- test_df %>% as_duckplyr_df_impl() %>% slice_head()
-  post <- test_df %>% slice_head() %>% as_duckplyr_df_impl()
+  pre <- test_df %>% as_duckplyr_df_impl() %>% slice_head(n = 2)
+  post <- test_df %>% slice_head(n = 2) %>% as_duckplyr_df_impl()
 
   # Compare
   expect_identical(pre, post)

tests/testthat/test-as_duckplyr_df.R

@@ -15007,6 +15007,57 @@ test_that("relational setdiff() order-enforcing", {
   DBI::dbDisconnect(con, shutdown = TRUE)
 })
 
+# slice_head order-preserving ----------------------------------------------------------
+
+test_that("relational slice_head(n = 2) order-preserving", {
+  # Autogenerated
+  duckdb <- asNamespace("duckdb")
+  drv <- duckdb::duckdb()
+  con <- DBI::dbConnect(drv)
+  experimental <- FALSE
+  df1 <- data.frame(a = seq(1, 6, by = 1), b = 2, g = c(1L, 2L, 2L, 3L, 3L, 3L))
+
+  "slice_head"
+  rel1 <- duckdb$rel_from_df(con, df1, experimental = experimental)
+  "slice_head"
+  rel2 <- duckdb$rel_limit(rel1, 2)
+  rel2
+  out <- duckdb$rel_to_altrep(rel2)
+  expect_identical(
+    out,
+    data.frame(a = c(1, 2), b = 2, g = 1:2)
+  )
+  DBI::dbDisconnect(con, shutdown = TRUE)
+})
+
+# slice_head order-enforcing -----------------------------------------------------------
+
+test_that("relational slice_head(n = 2) order-enforcing", {
+  # Autogenerated
+  duckdb <- asNamespace("duckdb")
+  drv <- duckdb::duckdb()
+  con <- DBI::dbConnect(drv)
+  experimental <- FALSE
+  df1 <- data.frame(a = seq(1, 6, by = 1), b = 2, g = c(1L, 2L, 2L, 3L, 3L, 3L))
+
+  "slice_head"
+  rel1 <- duckdb$rel_from_df(con, df1, experimental = experimental)
+  "slice_head"
+  rel2 <- duckdb$rel_limit(rel1, 2)
+  "arrange"
+  rel3 <- duckdb$rel_order(
+    rel2,
+    list(duckdb$expr_reference("a"), duckdb$expr_reference("b"), duckdb$expr_reference("g"))
+  )
+  rel3
+  out <- duckdb$rel_to_altrep(rel3)
+  expect_identical(
+    out,
+    data.frame(a = c(1, 2), b = 2, g = 1:2)
+  )
+  DBI::dbDisconnect(con, shutdown = TRUE)
+})
+
 # summarise order-preserving -----------------------------------------------------------
 
 test_that("relational summarise(c = mean(a)) order-preserving", {

tests/testthat/test-rel_api.R

@@ -881,6 +881,7 @@ test_extra_arg_map <- list(
     NULL
   ),
   semi_join = "join_by(a)",
+  slice_head = 'n = 2',
   slice_max = 'a',
   slice_min = 'a',
   summarise = c(
@@ -934,7 +935,6 @@ test_skip_map <- c(
   rowwise = "Stack overflow",
   sample_frac = "Random seed",
   sample_n = "Random seed",
-  slice_head = "External vector?",
   slice_max = "External vector?",
   slice_min = "External vector?",
   slice_sample = "External vector?",