Implement vec_get()

R/slice.R

@@ -125,6 +125,11 @@ vec_assign_fallback <- function(x, i, value) {
   x
 }
 
+vec_get <- function(x, i) {
+  i <- vec_as_position(i, vec_size(x), vec_names(x))
+  .Call(vctrs_get, x, i)
+}
+
 #' Create an index vector or a position
 #'
 #' @description

src/init.c

@@ -45,6 +45,7 @@ extern SEXP vctrs_typeof2(SEXP, SEXP);
 extern SEXP vctrs_cast(SEXP, SEXP, SEXP, SEXP);
 extern SEXP vctrs_as_index(SEXP, SEXP, SEXP, SEXP);
 extern SEXP vctrs_slice(SEXP, SEXP);
+extern SEXP vctrs_get(SEXP, SEXP);
 extern SEXP vctrs_init(SEXP, SEXP);
 extern SEXP vctrs_chop(SEXP, SEXP);
 extern SEXP vec_slice_seq(SEXP, SEXP, SEXP, SEXP);
@@ -137,6 +138,7 @@ static const R_CallMethodDef CallEntries[] = {
   {"vctrs_cast",                       (DL_FUNC) &vctrs_cast, 4},
   {"vctrs_as_index",                   (DL_FUNC) &vctrs_as_index, 4},
   {"vctrs_slice",                      (DL_FUNC) &vctrs_slice, 2},
+  {"vctrs_get",                        (DL_FUNC) &vctrs_get, 2},
   {"vctrs_init",                       (DL_FUNC) &vctrs_init, 2},
   {"vctrs_chop",                       (DL_FUNC) &vctrs_chop, 2},
   {"vctrs_slice_seq",                  (DL_FUNC) &vec_slice_seq, 4},

src/slice.c

@@ -206,14 +206,6 @@ static SEXP df_slice(SEXP x, SEXP index) {
     SET_VECTOR_ELT(out, i, sliced);
   }
 
-  SEXP row_nms = PROTECT(get_rownames(x));
-  if (TYPEOF(row_nms) == STRSXP) {
-    row_nms = PROTECT(slice_rownames(row_nms, index));
-    Rf_setAttrib(out, R_RowNamesSymbol, row_nms);
-    UNPROTECT(1);
-  }
-  UNPROTECT(1);
-
   UNPROTECT(1);
   return out;
 }
@@ -393,7 +385,15 @@ SEXP vec_slice_impl(SEXP x, SEXP index) {
 
   case vctrs_type_dataframe: {
     SEXP out = PROTECT_N(df_slice(data, index), &nprot);
+
+    SEXP row_names = PROTECT_N(get_rownames(data), &nprot);
+    if (TYPEOF(row_names) == STRSXP) {
+      row_names = PROTECT_N(slice_rownames(row_names, index), &nprot);
+      Rf_setAttrib(out, R_RowNamesSymbol, row_names);
+    }
+
     out = vec_restore(out, x, restore_size);
+
     UNPROTECT(nprot);
     return out;
   }
@@ -686,6 +686,140 @@ SEXP vctrs_as_index(SEXP i, SEXP n, SEXP names, SEXP convert_negative) {
 
 // -----------------------------------------------------------------------------
 
+static SEXP list_get(SEXP x, SEXP index) {
+  int i = INTEGER(index)[0];
+  return VECTOR_ELT(x, i - 1);
+}
+
+static SEXP get_shape_names(SEXP x) {
+  SEXP names = PROTECT(Rf_getAttrib(x, R_DimNamesSymbol));
+
+  if (names == R_NilValue) {
+    UNPROTECT(1);
+    return names;
+  }
+
+  names = PROTECT(Rf_shallow_duplicate(names));
+  SET_VECTOR_ELT(names, 0, R_NilValue);
+
+  UNPROTECT(2);
+  return names;
+}
+
+static SEXP vec_get_impl(SEXP x, SEXP index) {
+  int nprot = 0;
+
+  SEXP restore_size = PROTECT_N(r_int(1), &nprot);
+
+  struct vctrs_proxy_info info = vec_proxy_info(x);
+  PROTECT_PROXY_INFO(&info, &nprot);
+
+  SEXP data = info.proxy;
+
+  // Fallback to `[[` if the class doesn't implement a proxy. This is
+  // to be maximally compatible with existing classes.
+  if (vec_requires_fallback(x, info)) {
+    if (info.type == vctrs_type_scalar) {
+      Rf_errorcall(R_NilValue, "Can't extract from a scalar");
+    }
+
+    SEXP out;
+
+    if (has_dim(x)) {
+      out = PROTECT_N(vec_slice_fallback(x, index), &nprot);
+      Rf_setAttrib(out, R_DimNamesSymbol, get_shape_names(x));
+    } else {
+      out = PROTECT_N(
+        vctrs_dispatch2(syms_bracket_bracket, fns_bracket_bracket, syms_x, x, syms_i, index),
+        &nprot
+      );
+    }
+
+    // Take over attribute restoration only if the `[[` method did not
+    // restore itself
+    if (ATTRIB(out) == R_NilValue) {
+      out = vec_restore(out, x, restore_size);
+    }
+
+    UNPROTECT(nprot);
+    return out;
+  }
+
+  switch (info.type) {
+  case vctrs_type_null: {
+    Rf_error("Internal error: Unexpected `NULL` in `vec_get_impl()`.");
+  }
+  case vctrs_type_logical:
+  case vctrs_type_integer:
+  case vctrs_type_double:
+  case vctrs_type_complex:
+  case vctrs_type_character:
+  case vctrs_type_raw: {
+    SEXP out;
+
+    if (has_dim(x)) {
+      out = PROTECT_N(vec_slice_shaped(info.type, data, index), &nprot);
+      Rf_setAttrib(out, R_DimNamesSymbol, get_shape_names(data));
+    } else {
+      out = PROTECT_N(vec_slice_base(info.type, data, index), &nprot);
+    }
+
+    out = vec_restore(out, x, restore_size);
+
+    UNPROTECT(nprot);
+    return out;
+  }
+  case vctrs_type_list: {
+    SEXP out;
+
+    if (has_dim(x)) {
+      out = PROTECT_N(vec_slice_shaped(info.type, data, index), &nprot);
+      Rf_setAttrib(out, R_DimNamesSymbol, get_shape_names(data));
+
+      out = vec_restore(out, x, restore_size);
+
+      UNPROTECT(nprot);
+      return out;
+    }
+
+    out = list_get(data, index);
+
+    UNPROTECT(nprot);
+    return out;
+  }
+  case vctrs_type_dataframe: {
+    SEXP out = PROTECT_N(df_slice(data, index), &nprot);
+    out = vec_restore(out, x, restore_size);
+    UNPROTECT(nprot);
+    return out;
+  }
+  default:
+    Rf_error(
+      "Internal error: Unexpected type `%s` for vector proxy in `vec_get()`",
+      vec_type_as_str(info.type)
+    );
+  }
+}
+
+SEXP vec_get(SEXP x, SEXP index) {
+  vec_assert(x, args_empty);
+
+  // TODO - Currently using R level `vec_as_position()`
+  //index = PROTECT(vec_as_position(index, vec_size(x), PROTECT(vec_names(x))));
+
+  SEXP out = PROTECT(vec_get_impl(x, index));
+
+  UNPROTECT(1);
+  return out;
+}
+
+// [[ register() ]]
+SEXP vctrs_get(SEXP x, SEXP index) {
+  return vec_get(x, index);
+}
+
+// -----------------------------------------------------------------------------
+
 /*
  * @member proxy_info The result of `vec_proxy_info(x)`.
  * @member restore_size The restore size used in each call to `vec_restore()`.

src/utils.c

@@ -995,6 +995,7 @@ SEXP syms_y = NULL;
 SEXP syms_to = NULL;
 SEXP syms_dots = NULL;
 SEXP syms_bracket = NULL;
+SEXP syms_bracket_bracket = NULL;
 SEXP syms_x_arg = NULL;
 SEXP syms_y_arg = NULL;
 SEXP syms_to_arg = NULL;
@@ -1011,6 +1012,7 @@ SEXP syms_missing = NULL;
 SEXP syms_size = NULL;
 
 SEXP fns_bracket = NULL;
+SEXP fns_bracket_bracket = NULL;
 SEXP fns_quote = NULL;
 SEXP fns_names = NULL;
 
@@ -1174,6 +1176,7 @@ void vctrs_init_utils(SEXP ns) {
   syms_to = Rf_install("to");
   syms_dots = Rf_install("...");
   syms_bracket = Rf_install("[");
+  syms_bracket_bracket = Rf_install("[[");
   syms_x_arg = Rf_install("x_arg");
   syms_y_arg = Rf_install("y_arg");
   syms_to_arg = Rf_install("to_arg");
@@ -1190,6 +1193,7 @@ void vctrs_init_utils(SEXP ns) {
   syms_size = Rf_install("size");
 
   fns_bracket = Rf_findVar(syms_bracket, R_BaseEnv);
+  fns_bracket_bracket = Rf_findVar(syms_bracket_bracket, R_BaseEnv);
   fns_quote = Rf_findVar(Rf_install("quote"), R_BaseEnv);
   fns_names = Rf_findVar(Rf_install("names"), R_BaseEnv);
 

src/utils.h

@@ -242,6 +242,7 @@ extern SEXP syms_y;
 extern SEXP syms_to;
 extern SEXP syms_dots;
 extern SEXP syms_bracket;
+extern SEXP syms_bracket_bracket;
 extern SEXP syms_x_arg;
 extern SEXP syms_y_arg;
 extern SEXP syms_to_arg;
@@ -260,6 +261,7 @@ extern SEXP syms_size;
 #define syms_names R_NamesSymbol
 
 extern SEXP fns_bracket;
+extern SEXP fns_bracket_bracket;
 extern SEXP fns_quote;
 extern SEXP fns_names;
 

tests/testthat/test-slice.R

@@ -404,6 +404,144 @@ test_that("vec_slice() asserts vectorness (#301)", {
 })
 
 
+# vec_get -----------------------------------------------------------------
+
+test_that("vec_get throws error with non-vector inputs", {
+  expect_error(vec_get(environment(), 1L), class = "vctrs_error_scalar_type")
+})
+
+test_that("cannot extract more than one element", {
+  expect_error(vec_get(1, 1:2), class = "vctrs_error_position_bad_type")
+})
+
+test_that("can extract from atomic vectors", {
+  i <- 2L
+  expect_identical(vec_get(lgl(1, 0, 1), i), lgl(0))
+  expect_identical(vec_get(int(1, 2, 3), i), int(2))
+  expect_identical(vec_get(dbl(1, 2, 3), i), dbl(2))
+  expect_identical(vec_get(cpl(1, 2, 3), i), cpl(2))
+  expect_identical(vec_get(chr("1", "2", "3"), i), chr("2"))
+  expect_identical(vec_get(bytes(1, 2, 3), i), bytes(2))
+})
+
+test_that("can extract from a list", {
+  expect_identical(vec_get(list(1, 2, 3), 2L), 2)
+})
+
+test_that("can extract from shaped atomic vectors", {
+  i <- 2L
+  mat <- as.matrix
+  expect_identical(vec_get(mat(lgl(1, 0, 1)), i), mat(lgl(0)))
+  expect_identical(vec_get(mat(int(1, 2, 3)), i), mat(int(2)))
+  expect_identical(vec_get(mat(dbl(1, 2, 3)), i), mat(dbl(2)))
+  expect_identical(vec_get(mat(cpl(1, 2, 3)), i), mat(cpl(2)))
+  expect_identical(vec_get(mat(chr("1", "2", "3")), i), mat(chr("2")))
+  expect_identical(vec_get(mat(bytes(1, 2, 3)), i), mat(bytes(2)))
+})
+
+test_that("can extract from a shaped list", {
+  x <- matrix(list(1, 2, 3, 4), nrow = 2)
+  expect <- matrix(list(2, 4), nrow = 1)
+  expect_identical(vec_get(x, 2), expect)
+})
+
+test_that("can extract object of any dimensionality", {
+  x0 <- c(1, 1)
+  x1 <- ones(2)
+  x2 <- ones(2, 3)
+  x3 <- ones(2, 3, 4)
+  x4 <- ones(2, 3, 4, 5)
+
+  expect_equal(vec_get(x0, 1L), 1)
+  expect_identical(vec_get(x1, 1L), ones(1))
+  expect_identical(vec_get(x2, 1L), ones(1, 3))
+  expect_identical(vec_get(x3, 1L), ones(1, 3, 4))
+  expect_identical(vec_get(x4, 1L), ones(1, 3, 4, 5))
+})
+
+test_that("can extract from data frames row wise", {
+  df <- data.frame(x = 1:2, y = c("a", "b"))
+  expect_equal(vec_get(df, 1), vec_slice(df, 1))
+})
+
+test_that("can extract from data frames with data frame columns", {
+  df <- data.frame(x = 1:2)
+  df$y <- data.frame(a = 2:1)
+  expect_equal(vec_get(df, 1), vec_slice(df, 1))
+})
+
+test_that("names are lost from atomics", {
+  x <- set_names(1:2)
+  expect_equal(names(vec_get(x, 1)), NULL)
+})
+
+test_that("row names are lost from data frames", {
+  df <- data.frame(x = 1:2, row.names = c("r1", "r2"))
+  expect_equal(rownames(vec_get(df, 1)), "1")
+  expect_equal(.row_names_info(vec_get(df, 1)), -1)
+})
+
+test_that("row names are lost from matrices / arrays", {
+  x <- array(1, c(2, 2, 2), dimnames = list(c("r1", "r2")))
+  expect_equal(rownames(vec_get(x, 1)), NULL)
+})
+
+test_that("non-row names are kept on matrices / arrays", {
+  dim_names <- list(c("r1", "r2"), c("c1", "c2"), c("d1", "d2"))
+  x <- array(1, c(2, 2, 2), dimnames = dim_names)
+
+  expect <- dim_names
+  expect[1] <- list(NULL)
+
+  expect_equal(dimnames(vec_get(x, 1)), expect)
+})
+
+# TODO - Is this right?
+test_that("dimname names are kept on the 1st dimension", {
+  x <- array(1, c(2, 2), dimnames = list(kept = c("r1", "r2")))
+  expect_equal(dimnames(vec_get(x, 1)), list(kept = NULL, NULL))
+})
+
+test_that("row names are lost from shaped S3 objects", {
+  dim_names <- list(c("r1", "r2"), c("c1", "c2"))
+  x <- structure(1:4, dim = c(2L, 2L), dimnames = dim_names, class = "vctrs_mat")
+
+  expect_equal(dimnames(vec_get(x, 1)), list(NULL, c("c1", "c2")))
+})
+
+test_that("vec_get() falls back to `[[` with S3 objects", {
+  local_methods(
+    `[[.vctrs_foobar` = function(x, i, ...) "dispatched"
+  )
+  expect_identical(vec_get(foobar(NA), 1), foobar("dispatched"))
+})
+
+test_that("can extract from S3 lists that implement a proxy", {
+  expect_error(vec_get(foobar(list(NA)), 1), class = "vctrs_error_scalar_type")
+
+  local_methods(
+    vec_proxy.vctrs_foobar = identity
+  )
+
+  expect_identical(vec_get(foobar(list(NA)), 1), NA)
+})
+
+test_that("vec_get() doesn't restore when attributes have already been restored", {
+  local_methods(
+    `[[.vctrs_foobar` = function(x, i, ...) structure("dispatched", foo = "bar"),
+    vec_restore.vctrs_foobar = function(...) stop("not called")
+  )
+  expect_error(vec_get(foobar(NA), 1), NA)
+})
+
+test_that("vec_restore() is called after extracting from data frames", {
+  local_methods(
+    vec_restore.vctrs_tabble = function(...) "dispatched"
+  )
+  df <- structure(mtcars, class = c("vctrs_tabble", "data.frame"))
+  expect_identical(vec_get(df, 1), "dispatched")
+})
+
 # vec_init ----------------------------------------------------------------
 
 test_that("na of atomic vectors is as expected", {