updated logic

Author: venom1204

NEWS.md

@@ -48,6 +48,8 @@
 
 13. `setnafill()` now accepts a logical vector for the `cols` argument, which must be the same length as the number of columns in `x`, [#4113](https://github.com/Rdatatable/data.table/issues/4113). Thanks to @MichaelChirico for the suggestion and @venom1204 for the PR.
 
+14. `transpose()` and `tstrsplit()` gain a `keep` argument to specify which columns to return and in what order, [#5250](https://github.com/Rdatatable/data.table/issues/5250). When `keep` is used, the operation is now significantly more memory-efficient because memory for "throwaway" columns is never allocated in the C engine. Thanks to @MichaelChirico for the suggestion and @venom1204 the implementation.
+
 ### BUG FIXES
 
 1. `fread()` with `skip=0` and `(header=TRUE|FALSE)` no longer skips the first row when it has fewer fields than subsequent rows, [#7463](https://github.com/Rdatatable/data.table/issues/7463). Thanks @emayerhofer for the report and @ben-schwen for the fix.

R/transpose.R

@@ -1,4 +1,5 @@
-transpose = function(l, fill=NA, ignore.empty=FALSE, keep.names=NULL, make.names=NULL, list.cols=FALSE) {
+transpose = function(l, fill=NA, ignore.empty=FALSE, keep.names=NULL, make.names=NULL, list.cols=FALSE, keep=NULL) {
+  if (!is.null(keep)) keep = as.integer(keep)
   if (!is.null(make.names)) {
     stopifnot(length(make.names)==1L)
     if (is.character(make.names)) {
@@ -14,7 +15,7 @@ transpose = function(l, fill=NA, ignore.empty=FALSE, keep.names=NULL, make.names
     colnames = as.character(l[[make.names]])
     l = if (is.data.table(l)) l[,-make.names,with=FALSE] else l[-make.names]
   }
-  ans = .Call(Ctranspose, l, fill, ignore.empty, keep.names, list.cols)
+  ans = .Call(Ctranspose, l, fill, ignore.empty, keep.names, list.cols, keep)
   if (!is.null(make.names)) setattr(ans, "names", c(keep.names, colnames))
   else if (is.data.frame(l))  # including data.table but not plain list
     setattr(ans, "names", c(keep.names, paste0("V", seq_len(length(ans)-length(keep.names)))))
@@ -30,18 +31,22 @@ tstrsplit = function(x, ..., fill=NA, type.convert=FALSE, keep, names=FALSE, rev
     stopf("'rev' must be TRUE or FALSE.")
   ans = strsplit(as.character(x), ...)
   if (rev) ans = lapply(ans, base::rev)
-  ans = transpose(ans, fill=fill, ignore.empty=FALSE)
-
   if (!missing(keep)) {
     keep = suppressWarnings(as.integer(keep))
-    chk = min(keep) >= min(1L, length(ans)) & max(keep) <= length(ans)
-    if (!isTRUE(chk) || !length(keep))
-      stopf("'keep' should contain integer values between %d and %d.", min(1L, length(ans)), length(ans))
+    maxlen = if (length(ans)) max(vapply(ans, length, 0L)) else 0L
+    chk = min(keep, na.rm=TRUE) >= min(1L, maxlen) && max(keep, na.rm=TRUE) <= maxlen
+    if (!isTRUE(chk) || !length(keep) || anyNA(keep))
+      stopf("'keep' should contain integer values between %d and %d.", min(1L, maxlen), maxlen)
+    can_keep_early = isFALSE(type.convert)
+    ans = transpose(ans, fill=fill, ignore.empty=FALSE, keep = if (can_keep_early) keep else NULL)
   } else {
+    ans = transpose(ans, fill=fill, ignore.empty=FALSE)
     keep = seq_along(ans)
   }
-  if (isFALSE(type.convert))
-    ans = ans[keep]
+
+  if (isFALSE(type.convert)) {
+    if (length(ans) != length(keep)) ans = ans[keep]
+  }
   # Implementing #1094, but default FALSE
   else if (isTRUE(type.convert))
     ans = lapply(ans[keep], type.convert, as.is=TRUE)

inst/tests/tests.Rraw

@@ -21985,3 +21985,11 @@ DT3 = data.table(a=c(1,NA), b=c("x",NA), c=c(3,NA))
 test(2387.04, setnafill(copy(DT3), type="locf", cols=sapply(DT3, is.numeric)), data.table(a=c(1,1), b=c("x",NA), c=c(3,3)))
 test(2387.05, setnafill(copy(DT3), type="locf", cols=c(TRUE,NA,FALSE)), error="'cols' contains NA at position 2")
 test(2387.06, setnafill(copy(DT3), type="locf", cols=c(TRUE,FALSE)), error="'cols' is a logical vector of length 2 but there are 3 columns")
+
+# transpose() supports keep argument, #5250
+x = list(1:5, 6:10, 11:15)
+test(2388.01, transpose(x, keep=1:2), transpose(x)[1:2])
+test(2388.02, transpose(x, keep=c(5L,3L,1L)), transpose(x)[c(5,3,1)])
+test(2388.03, transpose(x, keep=c(1L,1L,3L)), transpose(x)[c(1,1,3)])
+test(2388.04, transpose(x, keep=NULL), transpose(x))
+test(2388.05, transpose(x, keep=1:2, fill=NA), transpose(x, fill=NA)[1:2])

man/transpose.Rd

@@ -7,7 +7,7 @@
 
 \usage{
 transpose(l, fill=NA, ignore.empty=FALSE, keep.names=NULL,
-          make.names=NULL, list.cols=FALSE)
+          make.names=NULL, list.cols=FALSE, keep=NULL)
 }
 \arguments{
   \item{l}{ A list, data.frame or data.table. }
@@ -16,6 +16,7 @@ transpose(l, fill=NA, ignore.empty=FALSE, keep.names=NULL,
   \item{keep.names}{The name of the first column in the result containing the names of the input; e.g. \code{keep.names="rn"}. By default \code{NULL} and the names of the input are discarded.}
   \item{make.names}{The name or number of a column in the input to use as names of the output; e.g. \code{make.names="rn"}. By default \code{NULL} and default names are given to the output columns.}
   \item{list.cols}{Default is \code{FALSE}. \code{TRUE} will avoid promoting types and return columns of type \code{list} instead. \code{factor} will always be cast to \code{character}.}
+  \item{keep}{An integer vector of column indices to keep and return. The columns will be returned in the order specified. If \code{NULL} (default), all columns are returned. This is much more memory efficient than transposing the entire list and then subsetting the result.}
 }
 \details{
   The list elements (or columns of \code{data.frame}/\code{data.table}) should be all \code{atomic}. If list elements are of unequal lengths, the value provided in \code{fill} will be used so that the resulting list always has all elements of identical lengths. The class of input object is also preserved in the transposed result.
@@ -48,8 +49,9 @@ l = list(1:3, c("a", "b", "c"))
 lapply(seq(length(l[[1]])), function(x) lapply(l, `[[`, x))
 transpose(l, list.cols=TRUE)
 
-ll = list(nm=c('x', 'y'), 1:2, 3:4)
-transpose(ll, make.names="nm")
+ll = list(1:5, 6:10)
+transpose(ll, keep=1:2)
+transpose(ll, keep=c(3, 1))
 }
 \seealso{
   \code{\link{data.table}}, \code{\link{tstrsplit}}

src/data.table.h

@@ -445,7 +445,7 @@ SEXP lookup(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
 SEXP overlaps(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
 SEXP whichwrapper(SEXP, SEXP);
 SEXP shift(SEXP, SEXP, SEXP, SEXP);
-SEXP transpose(SEXP, SEXP, SEXP, SEXP, SEXP);
+SEXP transpose(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
 SEXP anyNA(SEXP, SEXP);
 SEXP setlevels(SEXP, SEXP, SEXP);
 SEXP rleid(SEXP, SEXP);

src/transpose.c

@@ -2,7 +2,7 @@
 #include <Rdefines.h>
 #include <time.h>
 
-SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listColsArg)
+SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listColsArg, SEXP keepArg)
 {
   int nprotect = 0;
   if (!isNewList(l))
@@ -22,6 +22,11 @@ SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listCo
     error(_("'%s' must be TRUE or FALSE"), "list.cols");
   const bool listCol = LOGICAL_RO(listColsArg)[0];
 
+  const bool use_keep = !isNull(keepArg);
+  if (use_keep && !isInteger(keepArg)) error(_("'keep' must be an integer vector."));
+  const int *keep_ptr = use_keep ? INTEGER_RO(keepArg) : NULL;
+  const int keep_len = use_keep ? LENGTH(keepArg) : 0;
+
   // preprocessing
   int maxlen = 0, zerolen = 0;
   SEXPTYPE maxtype = 0;
@@ -37,8 +42,16 @@ SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listCo
     if (type > maxtype) maxtype = type;
   }
   if (listCol) maxtype = VECSXP; // need to keep preprocessing for zerolen
+  if (use_keep) {
+    for (int i=0; i<keep_len; i++) {
+      if (keep_ptr[i] == NA_INTEGER || keep_ptr[i] < 1 || keep_ptr[i] > maxlen)
+        error(_("'keep' index %d is out of bounds [1, %d]"), keep_ptr[i], maxlen);
+    }
+  }
+  const int num_cols = use_keep ? keep_len : maxlen;
+
   fill = PROTECT(coerceVector(fill, maxtype)); nprotect++;
-  SEXP ans = PROTECT(allocVector(VECSXP, maxlen + rn)); nprotect++;
+  SEXP ans = PROTECT(allocVector(VECSXP, num_cols + rn)); nprotect++;
   const int anslen = (ignore) ? (ln - zerolen) : ln;
   if (rn) {
     SEXP tt;
@@ -48,7 +61,7 @@ SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listCo
       if (length(VECTOR_ELT(l, i))) SET_STRING_ELT(tt, j++, STRING_ELT(lNames, i));
     }
   }
-  for (int i = 0; i < maxlen; i++) {
+  for (int i = 0; i < num_cols; i++) {
     SET_VECTOR_ELT(ans, i + rn, allocVector(maxtype, anslen));
   }
   const SEXP *ansp = SEXPPTR_RO(ans);
@@ -63,34 +76,39 @@ SEXP transpose(SEXP l, SEXP fill, SEXP ignoreArg, SEXP keepNamesArg, SEXP listCo
     case LGLSXP: {
       const int *ili = LOGICAL_RO(li);
       const int ifill = LOGICAL_RO(fill)[0];
-      for (int j = 0; j < maxlen; j++) {
-        LOGICAL(ansp[j + rn])[k] = j < len ? ili[j] : ifill;
+      for (int j = 0; j < num_cols; j++) {
+        int idx = use_keep ? (keep_ptr[j] - 1) : j;
+        LOGICAL(ansp[j + rn])[k] = idx < len ? ili[idx] : ifill;
       }
     } break;
     case INTSXP: {
       const int *ili = INTEGER_RO(li);
       const int ifill = INTEGER_RO(fill)[0];
-      for (int j = 0; j < maxlen; j++) {
-        INTEGER(ansp[j + rn])[k] = j < len ? ili[j] : ifill;
+      for (int j = 0; j < num_cols; j++) {
+        int idx = use_keep ? (keep_ptr[j] - 1) : j;
+        INTEGER(ansp[j + rn])[k] = idx < len ? ili[idx] : ifill;
       }
     } break;
     case REALSXP: {
       const double *dli = REAL_RO(li);
       const double dfill = REAL_RO(fill)[0];
-      for (int j = 0; j < maxlen; j++) {
-        REAL(ansp[j + rn])[k] = j < len ? dli[j] : dfill;
+      for (int j = 0; j < num_cols; j++) {
+        int idx = use_keep ? (keep_ptr[j] - 1) : j;
+        REAL(ansp[j + rn])[k] = idx < len ? dli[idx] : dfill;
       }
     } break;
     case STRSXP: {
       const SEXP sfill = STRING_ELT(fill, 0);
-      for (int j = 0; j < maxlen; j++) {
-        SET_STRING_ELT(ansp[j + rn], k, j < len ? STRING_ELT(li, j) : sfill);
+      for (int j = 0; j < num_cols; j++) {
+        int idx = use_keep ? (keep_ptr[j] - 1) : j;
+        SET_STRING_ELT(ansp[j + rn], k, idx < len ? STRING_ELT(li, idx) : sfill);
       }
     } break;
     case VECSXP: {
       const SEXP vfill = VECTOR_ELT(fill, 0);
-      for (int j = 0; j < maxlen; j++) {
-        SET_VECTOR_ELT(ansp[j + rn], k, j < len ? VECTOR_ELT(li, j) : vfill);
+      for (int j = 0; j < num_cols; j++) {
+        int idx = use_keep ? (keep_ptr[j] - 1) : j;
+        SET_VECTOR_ELT(ansp[j + rn], k, idx < len ? VECTOR_ELT(li, idx) : vfill);
       }
     } break;
     default: