ARROW-11136: [R] Bindings for is.nan

jonkeane · nealrichardson · commit fe990af45c52 · 2021-01-12T20:42:55.000-08:00
Closes apache#9167 from jonkeane/ARROW-11136-is.nan Authored-by: Jonathan Keane <jkeane@gmail.com> Signed-off-by: Neal Richardson <neal.p.richardson@gmail.com>
diff --git a/r/NAMESPACE b/r/NAMESPACE
@@ -64,6 +64,8 @@ S3method(is.na,ChunkedArray)
 S3method(is.na,Expression)
 S3method(is.na,Scalar)
 S3method(is.na,array_expression)
+S3method(is.nan,Array)
+S3method(is.nan,ChunkedArray)
 S3method(length,Array)
 S3method(length,ChunkedArray)
 S3method(length,Scalar)
diff --git a/r/R/array.R b/r/R/array.R
@@ -247,6 +247,9 @@ length.Array <- function(x) x$length()
 #' @export
 is.na.Array <- function(x) call_function("is_null", x)
 
+#' @export
+is.nan.Array <- function(x) call_function("is_nan", x)
+
 #' @export
 as.vector.Array <- function(x, mode) x$as_vector()
 
diff --git a/r/R/chunked-array.R b/r/R/chunked-array.R
@@ -130,6 +130,9 @@ as.vector.ChunkedArray <- function(x, mode) x$as_vector()
 #' @export
 is.na.ChunkedArray <- function(x) call_function("is_null", x)
 
+#' @export
+is.nan.ChunkedArray <- function(x) call_function("is_nan", x)
+
 #' @export
 `[.ChunkedArray` <- filter_rows
 
diff --git a/r/R/expression.R b/r/R/expression.R
@@ -111,7 +111,8 @@ cast_array_expression <- function(x, to_type, safe = TRUE, ...) {
 
 .unary_function_map <- list(
   "!" = "invert",
-  "is.na" = "is_null"
+  "is.na" = "is_null",
+  "is.nan" = "is_nan"
 )
 
 .binary_function_map <- list(
diff --git a/r/tests/testthat/test-chunked-array.R b/r/tests/testthat/test-chunked-array.R
@@ -121,6 +121,20 @@ test_that("ChunkedArray handles NA", {
   expect_equal(as.vector(is.na(x)), c(is.na(data[[1]]), is.na(data[[2]]), is.na(data[[3]])))
 })
 
+test_that("ChunkedArray handles NaN", {
+  data <- list(as.numeric(1:10), c(NaN, 2:10), c(1:3, NaN, 5L))
+  x <- chunked_array(!!!data)
+
+  expect_equal(x$type, float64())
+  expect_equal(x$num_chunks, 3L)
+  expect_equal(length(x), 25L)
+  expect_equal(as.vector(x), c(1:10, c(NaN, 2:10), c(1:3, NaN, 5)))
+
+  chunks <- x$chunks
+  expect_equal(as.vector(is.nan(chunks[[2]])), is.nan(data[[2]]))
+  expect_equal(as.vector(is.nan(x)), c(is.nan(data[[1]]), is.nan(data[[2]]), is.nan(data[[3]])))
+})
+
 test_that("ChunkedArray supports logical vectors (ARROW-3341)", {
   # with NA
   data <- purrr::rerun(3, sample(c(TRUE, FALSE, NA), 100, replace = TRUE))
diff --git a/r/tests/testthat/test-dataset.R b/r/tests/testthat/test-dataset.R
@@ -44,7 +44,7 @@ df1 <- tibble(
 second_date <- lubridate::ymd_hms("2017-03-09 07:01:02")
 df2 <- tibble(
   int = 101:110,
-  dbl = as.numeric(51:60),
+  dbl = c(as.numeric(51:59), NaN),
   lgl = rep(c(TRUE, FALSE, NA, TRUE, FALSE), 2),
   chr = letters[10:1],
   fct = factor(LETTERS[10:1]),
@@ -471,6 +471,17 @@ test_that("filter() with is.na()", {
   )
 })
 
+test_that("filter() with is.nan()", {
+  ds <- open_dataset(dataset_dir, partitioning = schema(part = uint8()))
+  expect_equivalent(
+    ds %>%
+      select(part, dbl) %>%
+      filter(!is.nan(dbl), part == 2) %>%
+      collect(),
+    tibble(part = 2L, dbl = df2$dbl[!is.nan(df2$dbl)])
+  )
+})
+
 test_that("filter() with %in%", {
   ds <- open_dataset(dataset_dir, partitioning = schema(part = uint8()))
   expect_equivalent(

Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,8 @@ cast_array_expression <- function(x, to_type, safe = TRUE, ...) {`
`111`	`111`
`112`	`112`	`.unary_function_map <- list(`
`113`	`113`	`"!" = "invert",`
`114`		`- "is.na" = "is_null"`
	`114`	`+ "is.na" = "is_null",`
	`115`	`+ "is.nan" = "is_nan"`
`115`	`116`	`)`
`116`	`117`
`117`	`118`	`.binary_function_map <- list(`