feat(r): Improve vctr class integration

r/geoarrow/NAMESPACE

@@ -1,6 +1,9 @@
 # Generated by roxygen2: do not edit by hand
 
 S3method("[",geoarrow_vctr)
+S3method("[<-",geoarrow_vctr)
+S3method("[[<-",geoarrow_vctr)
+S3method(as.character,geoarrow_vctr)
 S3method(as_geoarrow_array,character)
 S3method(as_geoarrow_array,default)
 S3method(as_geoarrow_array,nanoarrow_array)
@@ -12,10 +15,20 @@ S3method(as_geoarrow_array_stream,default)
 S3method(as_geoarrow_array_stream,geoarrow_vctr)
 S3method(as_geoarrow_array_stream,nanoarrow_array_stream)
 S3method(as_nanoarrow_array,sfc)
+S3method(as_nanoarrow_array_extension,geoarrow_extension_spec)
 S3method(as_nanoarrow_array_stream,geoarrow_vctr)
+S3method(as_nanoarrow_schema,geoarrow_vctr)
+S3method(convert_array,geoarrow_vctr)
+S3method(convert_array,sfc)
+S3method(convert_array,wk_wkb)
+S3method(convert_array,wk_wkt)
+S3method(convert_array,wk_xy)
+S3method(convert_array_extension,geoarrow_extension_spec)
+S3method(format,geoarrow_vctr)
 S3method(infer_geoarrow_schema,default)
 S3method(infer_geoarrow_schema,nanoarrow_array)
 S3method(infer_geoarrow_schema,nanoarrow_array_stream)
+S3method(infer_nanoarrow_ptype_extension,geoarrow_extension_spec)
 S3method(infer_nanoarrow_schema,geoarrow_vctr)
 S3method(infer_nanoarrow_schema,sfc)
 S3method(infer_nanoarrow_schema,wk_wkb)
@@ -26,6 +39,7 @@ S3method(wk_handle,geoarrow_vctr)
 S3method(wk_is_geodesic,geoarrow_vctr)
 export(as_geoarrow_array)
 export(as_geoarrow_array_stream)
+export(as_geoarrow_vctr)
 export(geoarrow_handle)
 export(geoarrow_schema_parse)
 export(geoarrow_writer)
@@ -36,7 +50,12 @@ export(na_extension_large_wkt)
 export(na_extension_wkb)
 export(na_extension_wkt)
 importFrom(nanoarrow,as_nanoarrow_array)
+importFrom(nanoarrow,as_nanoarrow_array_extension)
 importFrom(nanoarrow,as_nanoarrow_array_stream)
+importFrom(nanoarrow,as_nanoarrow_schema)
+importFrom(nanoarrow,convert_array)
+importFrom(nanoarrow,convert_array_extension)
+importFrom(nanoarrow,infer_nanoarrow_ptype_extension)
 importFrom(nanoarrow,infer_nanoarrow_schema)
 importFrom(wk,wk_crs)
 importFrom(wk,wk_handle)

r/geoarrow/R/array.R

@@ -315,3 +315,14 @@ as_validity_buffer <- function(x) {
 
   list(null_count == null_count, buffer = x$buffers[[2]])
 }
+
+# This really needs a helper in nanoarrow, but for now, we need a way to drop
+# the extension type and convert storage only for testing
+force_array_storage <- function(array) {
+  schema <- infer_nanoarrow_schema(array)
+  schema$metadata[["ARROW:extension:name"]] <- NULL
+  array_shallow <- nanoarrow::nanoarrow_allocate_array()
+  nanoarrow::nanoarrow_pointer_export(array, array_shallow)
+  nanoarrow::nanoarrow_array_set_schema(array_shallow, schema, validate = FALSE)
+}
+

r/geoarrow/R/nanoarrow-compat.R

@@ -0,0 +1,42 @@
+
+register_geoarrow_extension <- function() {
+  for (ext_name in geoarrow_extension_name_all()) {
+    nanoarrow::register_nanoarrow_extension(
+      ext_name,
+      nanoarrow::nanoarrow_extension_spec(subclass = "geoarrow_extension_spec")
+    )
+  }
+}
+
+#' @importFrom nanoarrow infer_nanoarrow_ptype_extension
+#' @export
+infer_nanoarrow_ptype_extension.geoarrow_extension_spec <- function(extension_spec, x, ...) {
+  new_geoarrow_vctr(list(), x, integer())
+}
+
+#' @importFrom nanoarrow convert_array_extension
+#' @export
+convert_array_extension.geoarrow_extension_spec <- function(extension_spec,
+                                                            array, to, ...) {
+  # For the default, this will dispatch to convert_array.geoarrow_vctr().
+  # This gets called if to is a base R type (e.g., integer())
+  stop(
+    sprintf(
+      "Can't convert geoarrow extension array to object of class '%s'",
+      class(to)[1]
+    )
+  )
+}
+
+#' @importFrom nanoarrow as_nanoarrow_array_extension
+#' @export
+as_nanoarrow_array_extension.geoarrow_extension_spec <- function(
+    extension_spec, x, ..., schema = NULL) {
+  as_geoarrow_array(x, schema = schema)
+}
+
+#' @importFrom nanoarrow convert_array
+#' @export
+convert_array.geoarrow_vctr <- function(array, to, ...) {
+  as_geoarrow_vctr(array, schema = as_nanoarrow_schema(to))
+}

r/geoarrow/R/sf-compat.R

@@ -1,4 +1,16 @@
 
+# exported in zzz.R
+st_as_sfc.geoarrow_vctr <- function(x, ..., promote_multi = FALSE) {
+  sfc <- wk::wk_handle(x, wk::sfc_writer(promote_multi))
+  sf::st_set_crs(sfc, sf::st_crs(wk::wk_crs(x)))
+}
+
+#' @export
+convert_array.sfc <- function(array, to, ..., sfc_promote_multi = FALSE) {
+  vctr <- as_geoarrow_vctr(array)
+  st_as_sfc.geoarrow_vctr(vctr, promote_multi = sfc_promote_multi)
+}
+
 #' @importFrom nanoarrow infer_nanoarrow_schema
 #' @export
 infer_nanoarrow_schema.sfc <- function(x, ...) {

r/geoarrow/R/type.R

@@ -67,6 +67,12 @@ na_extension_geoarrow <- function(geometry_type, dimensions = "XY",
   na_extension_geoarrow_internal(type_id, crs = crs, edges = edges)
 }
 
+geoarrow_extension_name_all <- function() {
+  c("geoarrow.wkt", "geoarrow.wkb", "geoarrow.point", "geoarrow.linestring",
+    "geoarrow.polygon", "geoarrow.multipoint", "geoarrow.mutlilinestring",
+    "geoarrow.multipolygon")
+}
+
 #' Inspect a GeoArrow schema
 #'
 #' @param schema A [nanoarrow_schema][nanoarrow::as_nanoarrow_schema]

r/geoarrow/R/vctr.R

@@ -1,4 +1,17 @@
 
+#' GeoArrow encoded arrays as R vectors
+#'
+#' @param x An object that works with [as_geoarrow_array_stream()]. Most
+#'   spatial objects in R already work with this method.
+#' @param ... Passed to [as_geoarrow_array_stream()]
+#' @param schema An optional `schema` (e.g., [na_extension_geoarrow()]).
+#'
+#' @return A vctr of class 'geoarrow_vctr'
+#' @export
+#'
+#' @examples
+#' as_geoarrow_vctr("POINT (0 1)")
+#'
 as_geoarrow_vctr <- function(x, ..., schema = NULL) {
   if (inherits(x, "geoarrow_vctr") && is.null(schema)) {
     return(x)
@@ -28,16 +41,80 @@ new_geoarrow_vctr <- function(chunks, schema, indices = NULL) {
 `[.geoarrow_vctr` <- function(x, i) {
   attrs <- attributes(x)
   x <- NextMethod()
-  # Assert slice?
+
+  if (is.null(vctr_as_slice(x))) {
+    stop(
+      "Can't subset geoarrow_vctr with non-slice (e.g., only i:j indexing is supported)"
+    )
+  }
+
   attributes(x) <- attrs
   x
 }
 
+#' @export
+`[<-.geoarrow_vctr` <- function(x, i, value) {
+  stop("subset assignment for geoarrow_vctr is not supported")
+}
+
+#' @export
+`[[<-.geoarrow_vctr` <- function(x, i, value) {
+  stop("subset assignment for geoarrow_vctr is not supported")
+}
+
+#' @export
+format.geoarrow_vctr <- function(x, ..., width = NULL, digits = NULL) {
+  if (is.null(width)) {
+    width <- getOption("width", 100L)
+  }
+
+  width <- max(width, 20)
+
+  if (is.null(digits)) {
+    digits <- getOption("digits", 7L)
+  }
+
+  digits <- max(digits, 0)
+
+  formatted_array <- geoarrow_kernel_call_scalar(
+    "format_wkt",
+    x,
+    options = c(
+      max_element_size_bytes = width - 10L,
+      precision = digits
+    ),
+    n = length(attr(x, "chunks"))
+  )
+
+  formatted_chr <- nanoarrow::convert_array_stream(
+    formatted_array,
+    character(),
+    size = length(x)
+  )
+
+  sprintf("<%s>", formatted_chr)
+}
+
+# Because RStudio's viewer uses this, we want to use the potentially abbreviated
+# WKT from the format method
+#' @export
+as.character.geoarrow_vctr <- function(x, ...) {
+  format(x, ...)
+}
+
 #' @export
 infer_nanoarrow_schema.geoarrow_vctr <- function(x, ...) {
   attr(x, "schema")
 }
 
+# Because zero-length vctrs are R's way of communicating "type", implement
+# as_nanoarrow_schema() here so that it works in places that expect a type
+#' @importFrom nanoarrow as_nanoarrow_schema
+#' @export
+as_nanoarrow_schema.geoarrow_vctr <- function(x, ...) {
+  attr(x, "schema")
+}
+
 #' @export
 as_geoarrow_array_stream.geoarrow_vctr <- function(x, ..., schema = NULL) {
   as_nanoarrow_array_stream.geoarrow_vctr(x, ..., schema = NULL)

r/geoarrow/R/wk-compat.R

@@ -112,6 +112,60 @@ infer_nanoarrow_schema.wk_xy <- function(x, ...) {
   wk_geoarrow_schema(x, na_extension_geoarrow, "POINT", dimensions = dims)
 }
 
+#' @export
+convert_array.wk_wkt <- function(array, to, ...) {
+  schema <- infer_nanoarrow_schema(array)
+  geo <- geoarrow_schema_parse(schema)
+  vctr <- as_geoarrow_vctr(array)
+
+  if (geo$extension_name == "geoarrow.wkt") {
+    out <- wk::new_wk_wkt(convert_array(force_array_storage(array)))
+  } else {
+    out <- wk::wk_handle(vctr, wk::wkt_writer())
+  }
+
+  wk::wk_crs(out) <- wk::wk_crs_output(vctr, to)
+  wk::wk_is_geodesic_output(vctr, to)
+  out
+}
+
+#' @export
+convert_array.wk_wkb <- function(array, to, ...) {
+  schema <- infer_nanoarrow_schema(array)
+  geo <- geoarrow_schema_parse(schema)
+  vctr <- as_geoarrow_vctr(array)
+
+  if (geo$extension_name == "geoarrow.wkb") {
+    storage <- convert_array(force_array_storage(array))
+    # Comes back as a blob::blob
+    attributes(storage) <- NULL
+    out <- wk::new_wk_wkb(storage)
+  } else {
+    out <- wk::wk_handle(vctr, wk::wkb_writer())
+  }
+
+  wk::wk_crs(out) <- wk::wk_crs_output(vctr, to)
+  wk::wk_is_geodesic_output(vctr, to)
+  out
+}
+
+#' @export
+convert_array.wk_xy <- function(array, to, ...) {
+  schema <- infer_nanoarrow_schema(array)
+  geo <- geoarrow_schema_parse(schema)
+  vctr <- as_geoarrow_vctr(array)
+
+  if (geo$extension_name == "geoarrow.point") {
+    out <- wk::as_xy(convert_array(force_array_storage(array)))
+  } else {
+    out <- wk::wk_handle(vctr, wk::xy_writer())
+  }
+
+  wk::wk_crs(out) <- wk::wk_crs_output(vctr, to)
+  wk::wk_is_geodesic(out) <- wk::wk_is_geodesic_output(vctr, to)
+  wk::as_xy(out, dims = names(unclass(to)))
+}
+
 wk_geoarrow_schema <- function(x, type_constructor, ...) {
   if (inherits(x, c("nanoarrow_array", "nanoarrow_array_stream"))) {
     schema <- nanoarrow::infer_nanoarrow_schema(x)

r/geoarrow/R/zzz.R

@@ -0,0 +1,71 @@
+
+# nocov start
+.onLoad <- function(...) {
+  register_geoarrow_extension()
+
+  s3_register("sf::st_as_sfc", "geoarrow_vctr")
+}
+
+# From the `vctrs` package (this function is intended to be copied
+# without attribution or license requirements to avoid a hard dependency on
+# vctrs:
+# https://github.com/r-lib/vctrs/blob/c2a7710fe55e3a2249c4fdfe75bbccbafcf38804/R/register-s3.R#L25-L31
+s3_register <- function(generic, class, method = NULL) {
+  stopifnot(is.character(generic), length(generic) == 1)
+  stopifnot(is.character(class), length(class) == 1)
+
+  pieces <- strsplit(generic, "::")[[1]]
+  stopifnot(length(pieces) == 2)
+  package <- pieces[[1]]
+  generic <- pieces[[2]]
+
+  caller <- parent.frame()
+
+  get_method_env <- function() {
+    top <- topenv(caller)
+    if (isNamespace(top)) {
+      asNamespace(environmentName(top))
+    } else {
+      caller
+    }
+  }
+  get_method <- function(method, env) {
+    if (is.null(method)) {
+      get(paste0(generic, ".", class), envir = get_method_env())
+    } else {
+      method
+    }
+  }
+
+  register <- function(...) {
+    envir <- asNamespace(package)
+
+    # Refresh the method each time, it might have been updated by
+    # `devtools::load_all()`
+    method_fn <- get_method(method)
+    stopifnot(is.function(method_fn))
+
+
+    # Only register if generic can be accessed
+    if (exists(generic, envir)) {
+      registerS3method(generic, class, method_fn, envir = envir)
+    } else if (identical(Sys.getenv("NOT_CRAN"), "true")) {
+      warning(sprintf(
+        "Can't find generic `%s` in package %s to register S3 method.",
+        generic,
+        package
+      ))
+    }
+  }
+
+  # Always register hook in case package is later unloaded & reloaded
+  setHook(packageEvent(package, "onLoad"), register)
+
+  # Avoid registration failures during loading (pkgload or regular)
+  if (isNamespaceLoaded(package)) {
+    register()
+  }
+
+  invisible()
+}
+# nocov end