Generate LIMIT statement from head.tbl_lazy

Fixes tidyverse#1725

NAMESPACE

@@ -116,7 +116,7 @@ S3method(groups,data.frame)
 S3method(groups,grouped_df)
 S3method(groups,tbl_cube)
 S3method(groups,tbl_lazy)
-S3method(head,tbl_sql)
+S3method(head,tbl_lazy)
 S3method(inner_join,data.frame)
 S3method(inner_join,tbl_df)
 S3method(inner_join,tbl_lazy)
@@ -218,6 +218,7 @@ S3method(sql_build,op_base_remote)
 S3method(sql_build,op_distinct)
 S3method(sql_build,op_filter)
 S3method(sql_build,op_group_by)
+S3method(sql_build,op_head)
 S3method(sql_build,op_join)
 S3method(sql_build,op_mutate)
 S3method(sql_build,op_rename)

NEWS.md

@@ -242,10 +242,8 @@ If you have written a dplyr backend, you'll need to make some minor changes to y
   provided new methods in your backend, you'll need to rewrite.
 
 * `select_query()` gains a distinct argument which is used for generating
-  queries for `distinct()`. It loses the `offset` and `limits` arguments
-  which are no longer used because cross-database support is patch
-  (because in general it doesn't make sense to think about the order of the
-  rows in a query).
+  queries for `distinct()`. It loses the `offset` argument which was
+  never used (and hence never tested). 
 
 * `src_translate_env()` has been replaced by `sql_translate_env()` which
   should have methods for the connection object.

R/sql-build.R

@@ -98,6 +98,10 @@ sql_build.op_mutate <- function(op, con, ...) {
   )
 }
 
+#' @export
+sql_build.op_head <- function(op, con, ...) {
+  select_query(sql_build(op$x, con), limit = op$args$n)
+}
 
 #' @export
 sql_build.op_group_by <- function(op, con, ...) {

R/sql-generic.R

@@ -19,10 +19,13 @@ sql_select <- function(con, select, from, where = NULL, group_by = NULL,
 #' @export
 sql_select.default <- function(con, select, from, where = NULL,
                                group_by = NULL, having = NULL,
-                               order_by = NULL, distinct = FALSE, ...) {
-
-  out <- vector("list", 6)
-  names(out) <- c("select", "from", "where", "group_by", "having", "order_by")
+                               order_by = NULL,
+                               limit = NULL,
+                               distinct = FALSE,
+                               ...) {
+  out <- vector("list", 7)
+  names(out) <- c("select", "from", "where", "group_by", "having", "order_by",
+    "limit")
 
   assert_that(is.character(select), length(select) > 0L)
   out$select <- build_sql(
@@ -59,6 +62,11 @@ sql_select.default <- function(con, select, from, where = NULL,
       escape(order_by, collapse = ", ", con = con))
   }
 
+  if (!is.null(limit)) {
+    assert_that(is.numeric(limit), length(limit) == 1L)
+    out$limit <- build_sql("LIMIT ", limit, con = con)
+  }
+
   escape(unname(compact(out)), collapse = "\n", parens = FALSE, con = con)
 }
 

R/sql-query.R

@@ -10,13 +10,15 @@ select_query <- function(from,
                          group_by = character(),
                          having = character(),
                          order_by = character(),
+                         limit = NULL,
                          distinct = FALSE) {
 
   stopifnot(is.character(select))
   stopifnot(is.character(where))
   stopifnot(is.character(group_by))
   stopifnot(is.character(having))
   stopifnot(is.character(order_by))
+  stopifnot(is.null(limit) || (is.numeric(limit) && length(limit) == 1L))
   stopifnot(is.logical(distinct), length(distinct) == 1L)
 
   structure(
@@ -27,7 +29,8 @@ select_query <- function(from,
       group_by = group_by,
       having = having,
       order_by = order_by,
-      distinct = distinct
+      distinct = distinct,
+      limit = limit
     ),
     class = c("select_query", "query")
   )
@@ -43,6 +46,7 @@ print.select_query <- function(x, ...) {
   if (length(x$group_by)) cat("Group by: ", named_commas(x$group_by), "\n", sep = "")
   if (length(x$order_by)) cat("Order by: ", named_commas(x$order_by), "\n", sep = "")
   if (length(x$having))   cat("Having:   ", named_commas(x$having), "\n", sep = "")
+  if (length(x$limit))    cat("Limit:    ", x$limit, "\n", sep = "")
 }
 
 

R/sql-render.R

@@ -25,7 +25,8 @@ sql_render.select_query <- function(query, con = NULL, ..., root = FALSE) {
 
   sql_select(
     con, query$select, from, where = query$where, group_by = query$group_by,
-    having = query$having, order_by = query$order_by, distinct = query$distinct,
+    having = query$having, order_by = query$order_by, limit = query$limit,
+    distinct = query$distinct,
     ...
   )
 }

R/tbl-lazy.R

@@ -75,6 +75,11 @@ group_by_.tbl_lazy <- function(.data, ..., .dots, add = TRUE) {
   add_op_single("group_by", .data, dots = dots, args = list(add = add))
 }
 
+#' @export
+head.tbl_lazy <- function(x, n = 6L, ...) {
+  add_op_single("head", x, args = list(n = n))
+}
+
 #' @export
 ungroup.tbl_lazy <- function(x, ...) {
   add_op_single("ungroup", x)

R/tbl-sql.r

@@ -78,11 +78,6 @@ dim.tbl_sql <- function(x) {
   c(NA, length(op_vars(x$ops)))
 }
 
-#' @export
-head.tbl_sql <- function(x, n = 6L, ...) {
-  collect(x, n = n, warn_incomplete = FALSE)
-}
-
 #' @export
 tail.tbl_sql <- function(x, n = 6L, ...) {
   stop("tail() is not supported by sql sources", call. = FALSE)

tests/testthat/test-sql-build.R

@@ -139,6 +139,17 @@ test_that("distinct sets flagged", {
 })
 
 
+# head --------------------------------------------------------------------
+
+test_that("head limits rows", {
+  out <- lazy_frame(x = 1:100) %>%
+    head(10) %>%
+    sql_build()
+
+  expect_equal(out$limit, 10)
+})
+
+
 # joins -------------------------------------------------------------------
 
 test_that("join captures both tables", {

tests/testthat/test-sql-render.R

@@ -52,6 +52,15 @@ test_that("distinct adds DISTINCT suffix", {
   expect_equal(out, data_frame(x = 1))
 })
 
+test_that("head limits rows returned", {
+  out <- memdb_frame(x = 1:100) %>%
+    head(10) %>%
+    collect()
+
+  expect_equal(nrow(out), 10)
+})
+
+
 test_that("mutate overwrites previous variables", {
   df <- memdb_frame(x = 1:5) %>%
     mutate(x = x + 1) %>%