test subclasses

Author: paleolimbot

r/R/convert-array.R

@@ -142,7 +142,11 @@ convert_fallback_other <- function(array, offset, length, to) {
 #' @export
 convert_array.nanoarrow_vctr <- function(array, to, ...) {
   schema <- attr(to, "schema", exact = TRUE)
-  as_nanoarrow_vctr(array, schema = schema)
+  if (is.null(schema)) {
+    schema <- infer_nanoarrow_schema(array)
+  }
+
+  new_nanoarrow_vctr(list(array), schema, class(to))
 }
 
 #' @export

r/R/vctr.R

@@ -17,8 +17,9 @@
 
 #' Experimental Arrow encoded arrays as R vectors
 #'
-#' @param x An object that works with [as_nanoarrow_array_stream()]. Most
-#'   spatial objects in R already work with this method.
+#' @param x An object that works with [as_nanoarrow_array_stream()].
+#' @param subclass An optional subclass of nanoarrow_vctr to prepend to the
+#'   final class name.
 #' @param ... Passed to [as_nanoarrow_array_stream()]
 #' @param schema An optional `schema`
 #'
@@ -29,14 +30,14 @@
 #' array <- as_nanoarrow_array(1:5)
 #' as_nanoarrow_vctr(array)
 #'
-as_nanoarrow_vctr <- function(x, ..., schema = NULL) {
+as_nanoarrow_vctr <- function(x, ..., schema = NULL, subclass = character()) {
   if (inherits(x, "nanoarrow_vctr") && is.null(schema)) {
     return(x)
   }
 
   stream <- as_nanoarrow_array_stream(x, ..., schema = schema)
   chunks <- collect_array_stream(stream, validate = FALSE)
-  new_nanoarrow_vctr(chunks, stream$get_schema())
+  new_nanoarrow_vctr(chunks, stream$get_schema(), subclass)
 }
 
 #' @rdname as_nanoarrow_vctr
@@ -102,7 +103,12 @@ format.nanoarrow_vctr <- function(x, ...) {
 }
 
 size_stable_format <- function(x, ...) {
-  if (inherits(x, "data.frame")) {
+  if (inherits(x, "nanoarrow_vctr")) {
+    # Extension types could have a default convert that gives a nanoarrow_vctr.
+    # If this is the case, they should be returning a subclass with a format
+    # method that ensures we don't get here.
+    rep(sprintf("<%s[%d]>", class(x)[1], seq_along(x)))
+  } else if (inherits(x, "data.frame")) {
     cols <- lapply(x, size_stable_format, ...)
     cols <- Map(paste, names(x), cols, sep = ": ")
     rows <- do.call(paste, c(cols, list(sep = ", ")))
@@ -258,7 +264,13 @@ print.nanoarrow_vctr <- function(x, ...) {
   stream <- as_nanoarrow_array_stream(utils::head(x, n_values))
   converted_head <- convert_array_stream(stream)
 
-  print(converted_head)
+  if (inherits(converted_head, "nanoarrow_vctr")) {
+    converted_head <- format(converted_head)
+    print(converted_head, quote = FALSE)
+  } else {
+    print(converted_head)
+  }
+
   if (more_values >= 2) {
     cat(sprintf("...and %d more values\n", more_values))
   } else if (more_values >= 1) {

r/man/as_nanoarrow_vctr.Rd

@@ -216,6 +216,22 @@ test_that("batched convert to vector works for nanoarrow_vctr()", {
   )
 })
 
+test_that("batched convert to vector works for nanoarrow_vctr() keeps subclass", {
+  vctr_ptype <- nanoarrow_vctr(subclass = "some_subclass")
+
+  empty_stream <- basic_array_stream(list(), schema = na_string())
+  empty_vctr <- convert_array_stream(empty_stream, vctr_ptype)
+  expect_s3_class(empty_vctr, "some_subclass")
+
+  stream1 <- basic_array_stream(list(c("")))
+  vctr1 <- convert_array_stream(stream1, vctr_ptype)
+  expect_s3_class(vctr1, "some_subclass")
+
+  stream2 <- basic_array_stream(list(c(""), c("")))
+  vctr2 <- convert_array_stream(stream2, vctr_ptype)
+  expect_s3_class(vctr2, "some_subclass")
+})
+
 test_that("convert to vector works for struct-style vectors", {
   array <- as_nanoarrow_array(as.POSIXlt("2021-01-01", tz = "America/Halifax"))
   expect_identical(

r/tests/testthat/test-convert-array.R

@@ -114,3 +114,40 @@ test_that("as_nanoarrow_array() dispatches on registered extension spec", {
     "some_ext"
   )
 })
+
+test_that("extensions can infer a schema of a nanoarrow_vctr() subclass", {
+  register_nanoarrow_extension(
+    "some_ext",
+    nanoarrow_extension_spec(subclass = "vctr_spec_class")
+  )
+  on.exit(unregister_nanoarrow_extension("some_ext"))
+
+  infer_nanoarrow_ptype_extension.vctr_spec_class <- function(spec, x, ...) {
+    nanoarrow_vctr(subclass = "some_vctr_subclass")
+  }
+
+  s3_register(
+    "nanoarrow::infer_nanoarrow_ptype_extension",
+    "vctr_spec_class",
+    infer_nanoarrow_ptype_extension.vctr_spec_class
+  )
+
+  expect_identical(
+    infer_nanoarrow_ptype(na_extension(na_string(), "some_ext")),
+    nanoarrow_vctr(subclass = "some_vctr_subclass")
+  )
+
+  ext_array <- nanoarrow_extension_array(c("one", "two", "three"), "some_ext")
+  vctr <- convert_array(ext_array)
+  expect_s3_class(vctr, "some_vctr_subclass")
+
+  # Ensure that registering a default conversion that returns a nanoarrow_vctr
+  # does not result in infinite recursion when printing or formatting it.
+  # An extension that does this should provide these methods for the subclass
+  # they return.
+  expect_length(format(vctr), length(vctr))
+  expect_output(
+    expect_identical(print(vctr), vctr),
+    "some_vctr_subclass"
+  )
+})