Skip to content

Commit

Permalink
gather and subset ghead/gtail
Browse files Browse the repository at this point in the history
  • Loading branch information
@ben-schwen
ben-schwen committed Aug 22, 2021
1 parent 9604551 commit 27a829c
Showing 1 changed file with 31 additions and 221 deletions.
252 changes: 31 additions & 221 deletions src/gsumm.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1067,125 +1067,27 @@ SEXP gtail(SEXP x, SEXP valArg) {
for (int i=0; i<ngrp; ++i){
if (val > grpsize[i]) anslen += grpsize[i]; else anslen += val;
}
SEXP ans;
switch(TYPEOF(x)) {
case LGLSXP: {
const int *ix = LOGICAL(x);
ans = PROTECT(allocVector(LGLSXP, anslen));
int *ians = LOGICAL(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
ians[offset+(rev-1-j)] = ix[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case INTSXP: {
const int *ix = INTEGER(x);
ans = PROTECT(allocVector(INTSXP, anslen));
int *ians = INTEGER(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
ians[offset+(rev-1-j)] = ix[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case REALSXP: {
const double *dx = REAL(x);
ans = PROTECT(allocVector(REALSXP, anslen));
double *dans = REAL(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
dans[offset+(rev-1-j)] = dx[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case CPLXSXP: {
const Rcomplex *dx = COMPLEX(x);
ans = PROTECT(allocVector(CPLXSXP, anslen));
Rcomplex *dans = COMPLEX(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
dans[offset+(rev-1-j)] = dx[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case STRSXP: {
ans = PROTECT(allocVector(STRSXP, anslen));
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
SET_STRING_ELT(ans, offset+(rev-1-j), STRING_ELT(x, k));
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case VECSXP: {
ans = PROTECT(allocVector(VECSXP, anslen));
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
SET_VECTOR_ELT(ans, offset+(rev-1-j), VECTOR_ELT(x, k));
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;

SEXP idx = allocVector(INTSXP, anslen);
int *idxp = INTEGER(idx);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
const int rev = (val > thisgrpsize ? thisgrpsize : val);
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+thisgrpsize-2-j;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
idxp[offset+(rev-1-j)] = k+1;
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
break;
default:
error(_("Type '%s' not supported by GForce gtail. Either add the prefix utils::tail(.) or turn off GForce optimization using options(datatable.optimize=1)"), type2char(TYPEOF(x)));
}
SEXP ans = PROTECT(allocVector(TYPEOF(x), anslen));
copyMostAttrib(x, ans);
subsetVectorRaw(ans, x, idx, false);
UNPROTECT(1);
return(ans);

}

SEXP ghead(SEXP x, SEXP valArg) {
Expand All @@ -1201,116 +1103,24 @@ SEXP ghead(SEXP x, SEXP valArg) {
const int thisgrpsize = grpsize[i];
if (val > thisgrpsize) anslen += thisgrpsize; else anslen += val;
}
SEXP ans;
switch(TYPEOF(x)) {
case LGLSXP: {
const int *ix = LOGICAL(x);
ans = PROTECT(allocVector(LGLSXP, anslen));
int *ians = LOGICAL(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
ians[offset+j] = ix[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case INTSXP: {
const int *ix = INTEGER(x);
ans = PROTECT(allocVector(INTSXP, anslen));
int *ians = INTEGER(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
ians[offset+j] = ix[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case REALSXP: {
const double *dx = REAL(x);
ans = PROTECT(allocVector(REALSXP, anslen));
double *dans = REAL(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
dans[offset+j] = dx[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case CPLXSXP: {
const Rcomplex *dx = COMPLEX(x);
ans = PROTECT(allocVector(CPLXSXP, anslen));
Rcomplex *dans = COMPLEX(ans);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
dans[offset+j] = dx[k];
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case STRSXP: {
ans = PROTECT(allocVector(STRSXP, anslen));
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
SET_STRING_ELT(ans, offset+j, STRING_ELT(x, k));
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
}
break;
case VECSXP: {
ans = PROTECT(allocVector(VECSXP, anslen));
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
SET_VECTOR_ELT(ans, offset+j, VECTOR_ELT(x, k));
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;

SEXP idx = allocVector(INTSXP, anslen);
int *idxp = INTEGER(idx);
int offset = 0;
for (int i=0; i<ngrp; ++i) {
const int thisgrpsize = grpsize[i];
for (int j=0; j<val; ++j) {
if (j > thisgrpsize-1) break;
int k = ff[i]+j-1;
if (isunsorted) k = oo[k]-1;
k = (irowslen == -1) ? k : irows[k]-1;
idxp[offset+j] = k+1;
}
if (val > thisgrpsize) offset += thisgrpsize; else offset += val;
}
break;
default:
error(_("Type '%s' not supported by GForce ghead. Either add the prefix utils::head(.) or turn off GForce optimization using options(datatable.optimize=1)"), type2char(TYPEOF(x)));
}
SEXP ans = PROTECT(allocVector(TYPEOF(x), anslen));
copyMostAttrib(x, ans);
subsetVectorRaw(ans, x, idx, false);
UNPROTECT(1);
return(ans);
}
Expand Down

0 comments on commit 27a829c

Please sign in to comment.