ARROW-11338: [R] Bindings for quantile and median

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))

r/DESCRIPTION

@@ -30,6 +30,7 @@ Imports:
  purrr,
  R6,
  rlang,
+ stats,
  tidyselect,
  utils,
  vctrs

r/NAMESPACE

@@ -58,6 +58,7 @@ S3method(match_arrow,ArrowDatum)
 S3method(match_arrow,default)
 S3method(max,ArrowDatum)
 S3method(mean,ArrowDatum)
+S3method(median,ArrowDatum)
 S3method(min,ArrowDatum)
 S3method(names,Dataset)
 S3method(names,FeatherReader)
@@ -73,6 +74,7 @@ S3method(print,array_expression)
 S3method(print,arrow_dplyr_query)
 S3method(print,arrow_info)
 S3method(print,arrow_r_metadata)
+S3method(quantile,ArrowDatum)
 S3method(read_message,InputStream)
 S3method(read_message,MessageReader)
 S3method(read_message,default)
@@ -152,6 +154,7 @@ export(ParquetFileWriter)
 export(ParquetVersionType)
 export(ParquetWriterProperties)
 export(Partitioning)
+export(QuantileInterpolation)
 export(RandomAccessFile)
 export(ReadableFile)
 export(RecordBatchFileReader)
@@ -308,6 +311,8 @@ importFrom(rlang,seq2)
 importFrom(rlang,set_names)
 importFrom(rlang,syms)
 importFrom(rlang,warn)
+importFrom(stats,median)
+importFrom(stats,quantile)
 importFrom(tidyselect,contains)
 importFrom(tidyselect,ends_with)
 importFrom(tidyselect,eval_select)

r/R/arrow-package.R

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+#' @importFrom stats quantile median
 #' @importFrom R6 R6Class
 #' @importFrom purrr as_mapper map map2 map_chr map_dfr map_int map_lgl keep
 #' @importFrom assertthat assert_that is.string

r/R/compute.R

@@ -80,6 +80,49 @@ collect_arrays_from_dots <- function(dots) {
  ChunkedArray$create(!!!arrays)
 }
 
+#' @export
+quantile.ArrowDatum <- function(x,
+ probs = seq(0, 1, 0.25),
+ na.rm = FALSE,
+ type = 7,
+ interpolation = c("linear", "lower", "higher", "nearest", "midpoint"),
+ ...) {
+ if (inherits(x, "Scalar")) x <- Array$create(x)
+ assert_is(probs, c("numeric", "integer"))
+ assert_that(length(probs) > 0)
+ assert_that(all(probs >= 0 & probs <= 1))
+ if (!na.rm && x$null_count > 0) {
+ stop("Missing values not allowed if 'na.rm' is FALSE", call. = FALSE)
+ }
+ if (type != 7) {
+ stop(
+ "Argument `type` not supported in Arrow. To control the quantile ",
+ "interpolation algorithm, set argument `interpolation` to one of: ",
+ "\"linear\" (the default), \"lower\", \"higher\", \"nearest\", or ",
+ "\"midpoint\".",
+ call. = FALSE
+ )
+ }
+ interpolation <- QuantileInterpolation[[toupper(match.arg(interpolation))]]
+ out <- call_function("quantile", x, options = list(q = probs, interpolation = interpolation))
+ if (length(out) == 0) {
+ # When there are no non-missing values in the data, the Arrow quantile
+ # function returns an empty Array, but for consistency with the R quantile
+ # function, we want an Array of NA_real_ with the same length as probs
+ out <- Array$create(rep(NA_real_, length(probs)))
+ }
+ out
+}
+
+#' @export
+median.ArrowDatum <- function(x, na.rm = FALSE, ...) {
+ if (!na.rm && x$null_count > 0) {
+ Scalar$create(NA_real_)
+ } else {
+ Scalar$create(quantile(x, probs = 0.5, na.rm = TRUE, ...))
+ }
+}
+
 #' @export
 unique.ArrowDatum <- function(x, incomparables = FALSE, ...) {
  call_function("unique", x)

r/R/enums.R

@@ -128,3 +128,9 @@ ParquetVersionType <- enum("ParquetVersionType",
 MetadataVersion <- enum("MetadataVersion",
  V1 = 0L, V2 = 1L, V3 = 2L, V4 = 3L, V5 = 4L
 )
+
+#' @export
+#' @rdname enums
+QuantileInterpolation <- enum("QuantileInterpolation",
+ LINEAR = 0L, LOWER = 1L, HIGHER = 2L, NEAREST = 3L, MIDPOINT = 4L
+)

r/src/compute.cpp

@@ -179,6 +179,23 @@ std::shared_ptr<arrow::compute::FunctionOptions> make_compute_options(
  return out;
  }
 
+ if (func_name == "quantile") {
+ using Options = arrow::compute::QuantileOptions;
+ auto out = std::make_shared<Options>(Options::Defaults());
+ SEXP q = options["q"];
+ if (!Rf_isNull(q) && TYPEOF(q) == REALSXP) {
+ out->q = cpp11::as_cpp<std::vector<double>>(q);
+ }
+ SEXP interpolation = options["interpolation"];
+ if (!Rf_isNull(interpolation) && TYPEOF(interpolation) == INTSXP &&
+ XLENGTH(interpolation) == 1) {
+ out->interpolation =
+ cpp11::as_cpp<enum arrow::compute::QuantileOptions::Interpolation>(
+ interpolation);
+ }
+ return out;
+ }
+
  if (func_name == "is_in" || func_name == "index_in") {
  using Options = arrow::compute::SetLookupOptions;
  return std::make_shared<Options>(cpp11::as_cpp<arrow::Datum>(options["value_set"]),

r/tests/testthat/test-compute-aggregate.R

@@ -199,6 +199,113 @@ test_that("Edge cases", {
  }
 })
 
+test_that("quantile.Array and quantile.ChunkedArray", {
+ a <- Array$create(c(0, 1, 2, 3))
+ ca <- ChunkedArray$create(c(0, 1), c(2, 3))
+ probs <- c(0.49, 0.51)
+ for(ad in list(a, ca)) {
+ for (type in c(int32(), uint64(), float64())) {
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "linear"),
+ Array$create(c(1.47, 1.53))
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "lower"),
+ Array$create(c(1, 1))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "higher"),
+ Array$create(c(2, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "nearest"),
+ Array$create(c(1, 2))$cast(type)
+ )
+ expect_equal(
+ quantile(ad$cast(type), probs = probs, interpolation = "midpoint"),
+ Array$create(c(1.5, 1.5))
+ )
+ }
+ }
+})
+
+test_that("quantile and median NAs, edge cases, and exceptions", {
+ expect_equal(
+ quantile(Array$create(c(1, 2)), probs = c(0, 1)),
+ Array$create(c(1, 2))
+ )
+ expect_error(
+ quantile(Array$create(c(1, 2, NA))),
+ "Missing values not allowed if 'na.rm' is FALSE"
+ )
+ expect_equal(
+ quantile(Array$create(numeric(0))),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_equal(
+ quantile(Array$create(rep(NA_integer_, 3)), na.rm = TRUE),
+ Array$create(rep(NA_real_, 5))
+ )
+ expect_error(
+ median(Array$create(c(1, 2)), probs = c(.25, .75))
+ )
+ expect_equal(
+ median(Array$create(c(1, 2)), interpolation = "higher"),
+ Scalar$create(2)
+ )
+ expect_equal(
+ quantile(Scalar$create(0L)),
+ Array$create(rep(0, 5))
+ )
+ expect_equal(
+ median(Scalar$create(1L)),
+ Scalar$create(1)
+ )
+ expect_error(
+ quantile(Array$create(1:3), type = 9),
+ "not supported"
+ )
+})
+
+test_that("median.Array and median.ChunkedArray", {
+ expect_vector_equal(
+ median(input),
+ 1:4
+ )
+ expect_vector_equal(
+ median(input),
+ 1:5
+ )
+ expect_vector_equal(
+ median(input),
+ numeric(0)
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+ expect_vector_equal(
+ median(input, na.rm = FALSE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ c(1, 2, NA)
+ )
+ expect_vector_equal(
+ median(input, na.rm = TRUE),
+ NA_real_
+ )
+})
+
 test_that("unique.Array", {
  a <- Array$create(c(1, 4, 3, 1, 1, 3, 4))
  expect_equal(unique(a), Array$create(c(1, 4, 3)))