Add more type checking and only floor() if necessary

r/R/arrow-datum.R

@@ -299,18 +299,23 @@ head.ArrowDatum <- function(x, n = 6L, ...) {
   } else {
     n <- min(len, n)
   }
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
   if (n == len) {
     return(x)
   }
-  x$Slice(0, floor(n))
+  x$Slice(0, n)
 }
 
 #' @importFrom utils tail
 #' @export
 tail.ArrowDatum <- function(x, n = 6L, ...) {
   assert_is(n, c("numeric", "integer"))
-  n <- floor(n)
   assert_that(length(n) == 1)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
   len <- NROW(x)
   if (n < 0) {
     # tail(x, negative) means all but the first n rows

r/R/dataset-scan.R

@@ -135,19 +135,29 @@ names.Scanner <- function(x) names(x$schema)
 
 #' @export
 head.Scanner <- function(x, n = 6L, ...) {
+  assert_is(n, c("numeric", "integer"))
+  assert_that(length(n) == 1)
   # Negative n requires knowing nrow(x), which requires a scan itself
   assert_that(n >= 0)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
   dataset___Scanner__head(x, floor(n))
 }
 
 #' @export
 tail.Scanner <- function(x, n = 6L, ...) {
-  tail_from_batches(dataset___Scanner__ScanBatches(x), floor(n))$read_table()
+  tail_from_batches(dataset___Scanner__ScanBatches(x), n)$read_table()
 }
 
 tail_from_batches <- function(batches, n) {
+  assert_is(n, c("numeric", "integer"))
+  assert_that(length(n) == 1)
   # Negative n requires knowing nrow(x), which requires a scan itself
-  assert_that(n >= 0) # For now
+  assert_that(n >= 0)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
   result <- list()
   batch_num <- 0
   # Given a list of batches, iterate from the back

r/R/dplyr.R

@@ -215,13 +215,23 @@ as.data.frame.arrow_dplyr_query <- function(x, row.names = NULL, optional = FALS
 
 #' @export
 head.arrow_dplyr_query <- function(x, n = 6L, ...) {
-  x$head <- floor(n)
+  assert_is(n, c("numeric", "integer"))
+  assert_that(length(n) == 1)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
+  x$head <- n
   collapse.arrow_dplyr_query(x)
 }
 
 #' @export
 tail.arrow_dplyr_query <- function(x, n = 6L, ...) {
-  x$tail <- floor(n)
+  assert_is(n, c("numeric", "integer"))
+  assert_that(length(n) == 1)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
+  x$tail <- n
   collapse.arrow_dplyr_query(x)
 }
 

r/R/record-batch-reader.R

@@ -130,8 +130,13 @@ as.data.frame.RecordBatchReader <- function(x, row.names = NULL, optional = FALS
 
 #' @export
 head.RecordBatchReader <- function(x, n = 6L, ...) {
+  assert_is(n, c("numeric", "integer"))
+  assert_that(length(n) == 1)
   # Negative n requires knowing nrow(x), which requires consuming the whole RBR
   assert_that(n >= 0)
+  if (!is.integer(n)) {
+    n <- floor(n)
+  }
   RecordBatchReader__Head(x, n)
 }