Implement list_sizes() (#539)

* Additional function vec_sizes() that is to vec_size() what lengths() is to length()

* Rework `vec_sizes()` into `list_sizes()`

* NEWS bullet

* Tweak NEWS bullet as suggested by code review

Co-authored-by: DavisVaughan <[email protected]>

NAMESPACE

@@ -461,6 +461,7 @@ export(fields)
 export(is_list_of)
 export(is_partial)
 export(list_of)
+export(list_sizes)
 export(maybe_lossy_cast)
 export(n_fields)
 export(new_data_frame)

NEWS.md

@@ -1,6 +1,11 @@
 
 # vctrs (development version)
 
+* New `list_sizes()` for computing the size of every element in a list.
+  `list_sizes()` is to `vec_size()` as `lengths()` is to `length()`, except
+  that it only supports lists. Atomic vectors and data frames result in an
+  error.
+
 * `vec_c()` and `vec_unchop()` now fall back to `base::c()` for S4 objects if
   the object doesn't implement `vec_ptype2()` but sets an S4 `c()`
   method (#919).

R/size.R

@@ -12,6 +12,12 @@
 #' frame (even data frame and matrix columns) have the same size.
 #' `vec_size_common(...)` returns the common size of multiple vectors.
 #'
+#' `list_sizes()` returns an integer vector containing the size of each element
+#' of a list. It is nearly equivalent to, but faster than,
+#' `map_int(x, vec_size)`, with the exception that `list_sizes()` will
+#' error on non-list inputs, as defined by [vec_is_list()]. `list_sizes()` is
+#' to `vec_size()` as [lengths()] is to [length()].
+#'
 #' @seealso [vec_slice()] for a variation of `[` compatible with `vec_size()`,
 #'   and [vec_recycle()] to recycle vectors to common length.
 #' @section Invariants:
@@ -68,6 +74,8 @@
 #' vec_size_common(1:10, 1:10)
 #' vec_size_common(1:10, 1)
 #' vec_size_common(integer(), 1)
+#'
+#' list_sizes(list("a", 1:5, letters))
 vec_size <- function(x) {
   .Call(vctrs_size, x)
 }
@@ -78,6 +86,12 @@ vec_size_common <- function(..., .size = NULL, .absent = 0L) {
   .External2(vctrs_size_common, .size, .absent)
 }
 
+#' @rdname vec_size
+#' @export
+list_sizes <- function(x) {
+  .Call(vctrs_list_sizes, x)
+}
+
 #' @rdname vec_size
 #' @export
 vec_is_empty <- function(x) {

src/init.c

@@ -36,6 +36,7 @@ extern SEXP vctrs_compare(SEXP, SEXP, SEXP);
 extern SEXP vctrs_match(SEXP, SEXP, SEXP);
 extern SEXP vctrs_duplicated_any(SEXP);
 extern SEXP vctrs_size(SEXP);
+extern SEXP vctrs_list_sizes(SEXP);
 extern SEXP vec_dim(SEXP);
 extern SEXP vctrs_dim_n(SEXP);
 extern SEXP vctrs_is_unspecified(SEXP);
@@ -153,6 +154,7 @@ static const R_CallMethodDef CallEntries[] = {
   {"vctrs_group_rle",                  (DL_FUNC) &vctrs_group_rle, 1},
   {"vctrs_group_loc",                  (DL_FUNC) &vec_group_loc, 1},
   {"vctrs_size",                       (DL_FUNC) &vctrs_size, 1},
+  {"vctrs_list_sizes",                 (DL_FUNC) &vctrs_list_sizes, 1},
   {"vctrs_dim",                        (DL_FUNC) &vec_dim, 1},
   {"vctrs_dim_n",                      (DL_FUNC) &vctrs_dim_n, 1},
   {"vctrs_is_unspecified",             (DL_FUNC) &vctrs_is_unspecified, 1},

src/size.c

@@ -98,6 +98,30 @@ SEXP vctrs_size(SEXP x) {
   return Rf_ScalarInteger(vec_size(x));
 }
 
+SEXP list_sizes(SEXP x) {
+  if (!vec_is_list(x)) {
+    Rf_errorcall(R_NilValue, "`x` must be a list.");
+  }
+
+  R_len_t size = vec_size(x);
+
+  SEXP out = PROTECT(Rf_allocVector(INTSXP, size));
+  int* p_out = INTEGER(out);
+
+  for (R_len_t i = 0; i < size; ++i) {
+    SEXP elt = VECTOR_ELT(x, i);
+    p_out[i] = vec_size(elt);
+  }
+
+  UNPROTECT(1);
+  return out;
+}
+
+// [[ register() ]]
+SEXP vctrs_list_sizes(SEXP x) {
+  return list_sizes(x);
+}
+
 R_len_t df_rownames_size(SEXP x) {
   for (SEXP attr = ATTRIB(x); attr != R_NilValue; attr = CDR(attr)) {
     if (TAG(attr) != R_RowNamesSymbol) {

tests/testthat/test-size.R

@@ -98,6 +98,17 @@ test_that("provided size is cast to an integer", {
   expect_identical(vec_size_common(.size = 1), 1L)
 })
 
+# list_sizes --------------------------------------------------------------
+
+test_that("only lists are allowed", {
+  expect_error(list_sizes(mtcars), "must be a list")
+  expect_error(list_sizes(1), "must be a list")
+})
+
+test_that("computes element sizes", {
+  expect_identical(list_sizes(list(1, 1:3, c("a", "b"))), c(1L, 3L, 2L))
+})
+
 # sequences ---------------------------------------------------------------
 
 test_that("vec_seq_along returns size-0 output for size-0 input", {