diff --git a/R/slice.R b/R/slice.R index fd48d7005..2b3dba2d4 100644 --- a/R/slice.R +++ b/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 diff --git a/src/init.c b/src/init.c index 2c2444129..739aa0b76 100644 --- a/src/init.c +++ b/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}, diff --git a/src/slice.c b/src/slice.c index d3cd58453..f7baf10ac 100644 --- a/src/slice.c +++ b/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()`. diff --git a/src/utils.c b/src/utils.c index bc2335db4..eb5cd64a7 100644 --- a/src/utils.c +++ b/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); diff --git a/src/utils.h b/src/utils.h index fa8390f8f..99bfe5e12 100644 --- a/src/utils.h +++ b/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; diff --git a/tests/testthat/test-slice.R b/tests/testthat/test-slice.R index 8d03ae1b9..7119f7782 100644 --- a/tests/testthat/test-slice.R +++ b/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", {