From 8f0eed215f8ceb282212afcd5ac0083c98cbc1bd Mon Sep 17 00:00:00 2001 From: Benjamin Schwendinger Date: Tue, 25 Aug 2026 15:33:05 +0200 Subject: [PATCH] GForce-optimize first()/last() with n>1 Adds gforce_dynamic: C level in gsumm.c recording that a GForce result returns MIN(w, grpsize[g]) items per group rather than the usual fixed 1 (or grpsize[g] for gshift). gforce() uses it to correctly replicate result rows. --- NEWS.md | 2 + R/data.table.R | 38 ++++++++------- inst/tests/optimize.Rraw | 37 ++++++++++++++- inst/tests/tests.Rraw | 4 +- src/data.table.h | 5 +- src/gsumm.c | 99 ++++++++++++++++++++++++++++++++++------ src/init.c | 2 + 7 files changed, 146 insertions(+), 41 deletions(-) diff --git a/NEWS.md b/NEWS.md index 5c2dd834f0..1db2d95578 100644 --- a/NEWS.md +++ b/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. `first()` and `last()` with `n>1` are now GForce optimized (e.g. `DT[, first(x, n=3), by=grp]`), [#4239](https://github.com/Rdatatable/data.table/issues/4239). Also adds a new internal `gforce_dynamic` mechanism to track any GForce result which returns other than exactly 1 row per group so that results are correctly replicated. Thanks to @nbenn for the report and @ben-schwen and @mattdowle for 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. diff --git a/R/data.table.R b/R/data.table.R index 7ac2f69917..52bf50888a 100644 --- a/R/data.table.R +++ b/R/data.table.R @@ -2093,30 +2093,21 @@ replace_dot_alias = function(e) { } } - # adding ghead/gtail(n) support for n > 1 #5060 #523 - q3 = 0L - if (!is.symbol(jsub)) { - headTail_arg = function(q) { - if (length(q)==3L && length(q3 <- q[[3L]])==1L && is.numeric(q3) && - (q[[1L]]) %chin% c("ghead", "gtail") && q3!=1L) q3 - else 0L - } - if (jsub %iscall% "list"){ - q3 = max(sapply(jsub, headTail_arg)) - } else if (length(jsub)==3L) { - q3 = headTail_arg(jsub) - } - } - if (q3 > 0L) { - grplens = pmin.int(q3, len__) - g = lapply(g, rep.int, times=grplens) - } else if (.is_nrows(jsub)) { + if (.is_nrows(jsub)) { + # multi-n shift() is not marked gforce_dynamic (its per-item results + # need unpacking into separate columns below), so it needs this branch + # specifically. g = lapply(g, rep.int, times=len__) # unpack list of lists for nrows functions zip_items = function(ll) do.call(mapply, c(list(FUN = c), ll, SIMPLIFY=FALSE, USE.NAMES=FALSE)) if (all(vapply_1b(ans, is.list))) { ans = lapply(ans, zip_items) } + } else if (!is.null(dynlens <- attr(ans, "gforce_dynamic", exact=TRUE))) { + # ghead/gtail/gfirst/glast with n>1 and shift() with a single n: + # gforce() has already validated 'ans' is consistently 'dynlens' items + # per group in gsumm.c + g = lapply(g, rep.int, times=dynlens) } ans = c(g, ans) } else { @@ -3346,8 +3337,8 @@ gfuns = c(gdtfuns, `g[` = `g[[` = function(x, n) .Call(Cgnthvalue, x, as.integer(n)) # n is of length=1 here. ghead = function(x, n) .Call(Cghead, x, as.integer(n)) gtail = function(x, n) .Call(Cgtail, x, as.integer(n)) -gfirst = function(x) .Call(Cgfirst, x) -glast = function(x) .Call(Cglast, x) +gfirst = function(x, n=1L) .Call(Cgfirst, x, as.integer(n)) +glast = function(x, n=1L) .Call(Cglast, x, as.integer(n)) gsum = function(x, na.rm=FALSE) .Call(Cgsum, x, na.rm) gmean = function(x, na.rm=FALSE) .Call(Cgmean, x, na.rm) gweighted.mean = function(x, w, ..., na.rm=FALSE) { @@ -3399,6 +3390,12 @@ is_constantish = function(q, check_singleton=FALSE) { length(q) == 3L && is_constantish(q[[3L]], check_singleton = TRUE) } +# first(x, n) / last(x, n) with n>1, #4446 #4239. +.gfirstlast_ok = function(q, envir) { + length(q) == 3L && + is_constantish(q[[3L]], check_singleton = TRUE) && + is.numeric(n <- eval(q[[3L]], envir)) && length(n)==1L && !is.na(n) && n>=1 +} `.g[_ok` = function(q, x, envir=parent.frame(3L)) { length(q) == 3L && is_constantish(q[[3L]], check_singleton = TRUE) && @@ -3458,6 +3455,7 @@ is_constantish = function(q, check_singleton=FALSE) { "shift" = .gshift_ok(q), "weighted.mean" = .gweighted.mean_ok(q, x), "tail" = , "head" = .ghead_ok(q), + "first" = , "last" = .gfirstlast_ok(q, envir), "[[" = , "[" = `.g[_ok`(q, x, envir), FALSE )) diff --git a/inst/tests/optimize.Rraw b/inst/tests/optimize.Rraw index 4c25baee8b..6687961d8d 100644 --- a/inst/tests/optimize.Rraw +++ b/inst/tests/optimize.Rraw @@ -110,7 +110,7 @@ test(1184.2,optimize=1L, DT[, mean(v), by=x, verbose=TRUE], output="(GForce FALS test(1185.1,optimize=c(0L, 1L, 2L), DT[, list(sum(y), sum(v), sum(y,na.rm=TRUE), sum(v,na.rm=TRUE)), by=x], data.table(x=c("a","b","c","d"), V1=c(NA,10L,NA,NA), V2=c(6,NA,NA,NA), V3=c(4L,10L,7L,0L), V4=c(6,10,15,0))) # test 1185.1 subsumes 1185.1 and 1187.1 for testing all levels -test(1185.2,optimize=c(0L,1L,Inf), DT[, list(mean(y), mean(v), mean(y,na.rm=TRUE), mean(v,na.rm=TRUE)), by=x, verbose=TRUE], +test(1185.2,optimize=c(0L,1L,Inf), DT[, list(mean(y), mean(v), mean(y,na.rm=TRUE), mean(v,na.rm=TRUE)), by=x, verbose=TRUE], output=c("All optimizations.*off", "Old mean.*changed j", "GForce optimized j to")) # test 1185.2 subsumes 1185.2, 1185.3, and 1185.4 for testing all levels MyVar = TRUE @@ -334,7 +334,7 @@ test(1594.08,optimize=opt, DT[(idx), var(z1, na.rm=TRUE), by=x], ans) test(1594.09,optimize=opt,DT[, lapply(.SD, sd, na.rm=FALSE), by=x]) test(1594.10,optimize=opt, DT[, lapply(.SD, sd, na.rm=TRUE), by=x], DT[, lapply(.SD, stats::sd, na.rm=TRUE), by=x]) test(1594.11,optimize=opt, DT[, lapply(.SD, sd, na.rm=TRUE), by=x, verbose=TRUE], output=out) - + test(1594.12,optimize=opt, DT[, lapply(.SD, prod, na.rm=FALSE), by=x]) test(1594.13,optimize=opt, DT[, lapply(.SD, prod, na.rm=TRUE), by=x]) test(1594.14,optimize=opt, DT[, lapply(.SD, prod, na.rm=TRUE), by=x, verbose=TRUE], output=out) @@ -490,3 +490,36 @@ test(2283 + 0.19, optimize=opts, names(M[, c(mpg[1], list(mpg, b=hp), c(lapply(.SD, mean))), by="cyl", .SDcols=c("vs", "am")]), c("cyl", "V1", "V2", "b", "vs", "am"), context=sprintf("optimize=%s [XIX]", format(opt))) + +# first()/last() with n>1 are now GForce optimized via gfirst()/glast(), #4239 +opt = 0:2 +DT = data.table(g=c(1,1,1,2,2), x=1:5, y=11:15) +out = c("GForce FALSE", "GForce FALSE", "GForce optimized j to") +test(2284.01, optimize=opt, DT[, first(x, n=2), by=g, verbose=TRUE], + data.table(g=c(1,1,2,2), V1=c(1L,2L,4L,5L)), output=out) +test(2284.02, optimize=opt, DT[, .(first(x,n=2), mean(y)), by=g, verbose=TRUE], + data.table(g=c(1,1,2,2), V1=c(1L,2L,4L,5L), V2=c(12,12,14.5,14.5)), output=out) +test(2284.03, optimize=opt, DT[, .(first(x,n=2), first(y,n=3)), by=g], error="Supplied 2 items for column 1 of group 1 which has 3 rows") +DT = data.table(a=1:3,b=(1:9)/10) +test(2284.04, optimize=opt, copy(DT)[, v := first(b, n=3), a, verbose=TRUE], copy(DT)[, v := b], output=out) +test(2284.05, DT[, v := first(b, n=2), a], error="Supplied 6 items to be assigned to 9 items of column 'v'.") +DT = data.table(g=c(1,1,1,2,2), s=c("a","b","c","d","e"), l=list(1,2,3,4,5)) +test(2284.06, DT[, .(first(s,n=2), last(l,n=2)), by=g], + data.table(g=c(1,1,2,2), V1=c("a","b","d","e"), V2=list(2,3,4,5))) +DT4 = data.table(g=c(1,1,1,2,2), lg=c(TRUE,FALSE,TRUE,FALSE,TRUE), db=c(1.5,2.5,3.5,4.5,5.5), cx=complex(1:5, 5:1)) +test(2284.07, DT4[, .(first(lg,n=2), first(db,n=2), first(cx,n=2)), by=g], + data.table(g=c(1,1,2,2), V1=c(TRUE,FALSE,FALSE,TRUE), V2=c(1.5,2.5,4.5,5.5), + V3=complex(c(1,2,4,5), c(5,4,2,1)))) +# fall back for n<1L +DT = data.table(g=c(1,1,1,2,2), x=1:5, y=11:15) +test(2284.08, DT[, first(x, n=0L), by=g, verbose=TRUE], + data.table(g=numeric(), V1=integer()), output="GForce is on, but not activated") + +# ghead()/gtail() and gshift() now also go through gforce_dynamic too +DT5 = data.table(g=c(1,1,1,2,2), x=1:5, y=1:5) +test(2285.01, optimize=opt, DT5[, .(shift(x), mean(x)), by=g, verbose=TRUE], + data.table(g=c(1,1,1,2,2), V1=c(NA,1L,2L,NA,4L), V2=c(2,2,2,4.5,4.5)), output=out) +test(2285.02, copy(DT5)[, c("s","m") := .(shift(x), mean(x)), by=g], + data.table(g=c(1,1,1,2,2), x=1:5, y=1:5, s=c(NA,1L,2L,NA,4L), m=c(2,2,2,4.5,4.5))) +test(2285.03, optimize=opt, DT5[, .(head(x,2), head(y,3)), by=g], error="Supplied 2 items for column 1 of group 1 which has 3 rows") +# head/tail with := and mismatched n: previously errored (Supplied N items...), now aligns/pads instead, see tests 2233.28/2233.29 in tests.Rraw diff --git a/inst/tests/tests.Rraw b/inst/tests/tests.Rraw index 7cb84b7ec1..6ca93a93f8 100644 --- a/inst/tests/tests.Rraw +++ b/inst/tests/tests.Rraw @@ -17857,8 +17857,8 @@ test(2233.25, copy(DT)[a!=4, v := head(b, 3L), a, verbose=TRUE], copy(DT)[a!=4, DT = data.table(a=1:3,b=(1:9)/10) test(2233.26, DT[, c("v1","v2") := .(min(b), max(b)), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v1=(1:3)/10, v2=(7:9)/10), output="GForce optimized j to") test(2233.27, DT[, c("v1","v2") := .(head(b,3L), tail(b,3L)), a, verbose=TRUE], data.table(a=1:3, b=(1:9)/10, v1=(1:9)/10, v2=(1:9)/10), output="GForce optimized j to") -test(2233.28, DT[, c("v1","v2") := .(head(b,3L), tail(b,2L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v2'.") -test(2233.29, DT[, c("v1","v2") := .(head(b,2L), tail(b,3L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v1'.") +test(2233.28, DT[, c("v1","v2") := .(head(b,3L), tail(b,2L)), a], error="Supplied 2 items for column 2 of group 1 which has 3 rows") +test(2233.29, DT[, c("v1","v2") := .(head(b,2L), tail(b,3L)), a], error="Supplied 2 items for column 1 of group 1 which has 3 rows") test(2233.30, DT[, c("v1","v2") := .(head(b,2L), tail(b,2L)), a], error="Supplied 6 items to be assigned to 9 items of column 'v1'.") test(2233.31, DT[, c("v1","v2") := .(min(b), max(b)), a, verbose=TRUE], DT[, c("v1","v2") := .(base::min(b), base::max(b)), a ], output="GForce optimized j to") test(2233.32, DT[, c("v1","v2") := .(head(b,3L), tail(b,3L)), a, verbose=TRUE], DT[, c("v1","v2") := .(utils::head(b,3L), utils::tail(b,3L)), a], output="GForce optimized j to") diff --git a/src/data.table.h b/src/data.table.h index df46c4a33b..a78327423f 100644 --- a/src/data.table.h +++ b/src/data.table.h @@ -137,6 +137,7 @@ extern SEXP sym_index; extern SEXP sym_BY; extern SEXP sym_starts, char_starts; extern SEXP sym_maxgrpn; +extern SEXP sym_gforce_dynamic; extern SEXP sym_anyna; extern SEXP sym_anyinfnan; extern SEXP sym_anynotascii; @@ -452,8 +453,8 @@ SEXP rleid(SEXP, SEXP); SEXP gmedian(SEXP, SEXP); SEXP gtail(SEXP, SEXP); SEXP ghead(SEXP, SEXP); -SEXP glast(SEXP); -SEXP gfirst(SEXP); +SEXP glast(SEXP, SEXP); +SEXP gfirst(SEXP, SEXP); SEXP gnthvalue(SEXP, SEXP); SEXP dim(SEXP); SEXP warn_matrix_column_r(SEXP); diff --git a/src/gsumm.c b/src/gsumm.c index 5051d72ca8..886dcaa29c 100644 --- a/src/gsumm.c +++ b/src/gsumm.c @@ -37,6 +37,59 @@ static int nbit(int n) return nb; } +/* Some GForce functions (currently gfirst/glast/ghead/gtail with n>1, and gshift) return, for at least + one group, a different number of items than 1 (the norm for e.g. gmean, gsum). Such a column carries + a 'gforce_dynamic' attribute (set by gfirstlast and gshift, below): a scalar integer w, meaning the + column has MIN(w, grpsize[g]) items for group g (gshift sets w=INT_MAX, since it always returns + exactly grpsize[g]). When a query combines such a column with an ordinary (fixed, 1-per-group) result + in the same by= (e.g. .(shift(x), mean(y))), this function replicates the fixed result out to match, + #1414. When a query combines two dynamic columns that don't actually agree in length for some group + (e.g. .(head(x,2), head(y,3)) where the group has >=3 rows), this errors rather than reconciling them + with NA. Otherwise optimize>=2 (GForce) would succeed at a query that optimize<2 (dogroups.c) can't + do at all, making the optimize level change correctness rather than just speed Returns gans unchanged + (no allocation) when no column carries the attribute, so the (very common) case of no dynamic-length + GForce result in the query costs nothing extra. */ +static SEXP gforce_align_dynamic(SEXP gans) { + const int nans = length(gans); + int max_w = 0; + for (int i=0; imax_w) max_w=this_w; + } + if (!max_w) return gans; // nothing in this j is gforce_dynamic; most common case, return untouched with no allocation + int nprotect = 0; + SEXP lens = PROTECT(allocVector(INTSXP, ngrp)); nprotect++; + int *lensp = INTEGER(lens); + for (int g=0; g1 is present, #4446 #4239 + UNPROTECT(nprotect); return ans; } @@ -922,7 +976,10 @@ SEXP gmedian(SEXP x, SEXP narmArg) { static SEXP gfirstlast(SEXP x, const bool first, const int w, const bool headw) { // w: which item (1 other than for gnthvalue when could be >1) - // headw: select 1:w of each group when first=true, and (n-w+1):n when first=false (i.e. tail) + // headw: select 1:w of each group when first=true, and (n-w+1):n when first=false (i.e. tail). + // When TRUE, the result is marked with a 'gforce_dynamic' attribute #4239 + // so that gforce() can correctly replicate ordinary (fixed, 1-per-group) results + // against it, and validate it against other dynamic results, when combined in the same by=. const bool nosubset = irowslen == -1; const bool issorted = !isunsorted; // make a const-bool for use inside loops const int n = nosubset ? length(x) : irowslen; @@ -988,29 +1045,34 @@ static SEXP gfirstlast(SEXP x, const bool first, const int w, const bool headw) default: error(_("Type '%s' is not supported by GForce head/tail/first/last/`[`. Either add the namespace prefix (e.g. utils::head(.)) or turn off GForce optimization using options(datatable.optimize=1)"), type2char(TYPEOF(x))); } + if (headw) setAttrib(ans, sym_gforce_dynamic, ScalarInteger(w)); // so gforce() can recompute MIN(w, grpsize[g]) per group copyMostAttrib(x, ans); UNPROTECT(1); return(ans); } -SEXP glast(SEXP x) { - return gfirstlast(x, false, 1, false); +SEXP glast(SEXP x, SEXP nArg) { + if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "glast is only implemented for n>0. This should have been caught before"); // # nocov + const int n=INTEGER(nArg)[0]; + return gfirstlast(x, false, n, n>1); } -SEXP gfirst(SEXP x) { - return gfirstlast(x, true, 1, false); +SEXP gfirst(SEXP x, SEXP nArg) { + if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gfirst is only implemented for n>0. This should have been caught before"); // # nocov + const int n=INTEGER(nArg)[0]; + return gfirstlast(x, true, n, n>1); } SEXP gtail(SEXP x, SEXP nArg) { if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gtail is only implemented for n>0. This should have been caught before"); // # nocov const int n=INTEGER(nArg)[0]; - return n==1 ? glast(x) : gfirstlast(x, false, n, true); + return gfirstlast(x, false, n, n>1); } SEXP ghead(SEXP x, SEXP nArg) { - if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "gtail is only implemented for n>0. This should have been caught before"); // # nocov + if (!isInteger(nArg) || LENGTH(nArg)!=1 || INTEGER(nArg)[0]<1) internal_error(__func__, "ghead is only implemented for n>0. This should have been caught before"); // # nocov const int n=INTEGER(nArg)[0]; - return n==1 ? gfirst(x) : gfirstlast(x, true, n, true); + return gfirstlast(x, true, n, n>1); } SEXP gnthvalue(SEXP x, SEXP nArg) { @@ -1279,7 +1341,14 @@ SEXP gshift(SEXP x, SEXP nArg, SEXP fillArg, SEXP typeArg) { } copyMostAttrib(x, tmp); // needed for integer64 because without not the correct class of int64 is assigned } - UNPROTECT(nprotect); // consistency with plain shift(): "strip" the list in the 1-input case, for convenience - return isVectorAtomic(x) && length(ans) == 1 ? VECTOR_ELT(ans, 0) : ans; + if (isVectorAtomic(x) && length(ans) == 1) { + SEXP tmp = VECTOR_ELT(ans, 0); + // mark single gshift with gforce_dynamic to use MIN(w, grpsize[g]) + setAttrib(tmp, sym_gforce_dynamic, ScalarInteger(INT_MAX)); + UNPROTECT(nprotect); + return tmp; + } + UNPROTECT(nprotect); + return ans; } diff --git a/src/init.c b/src/init.c index ec1ff384cd..36cad4a443 100644 --- a/src/init.c +++ b/src/init.c @@ -29,6 +29,7 @@ SEXP sym_index; SEXP sym_BY; SEXP sym_starts, char_starts; SEXP sym_maxgrpn; +SEXP sym_gforce_dynamic; SEXP sym_anyna; SEXP sym_anyinfnan; SEXP sym_anynotascii; @@ -336,6 +337,7 @@ void attribute_visible R_init_data_table(DllInfo *info) sym_index = install("index"); sym_BY = install(".BY"); sym_maxgrpn = install("maxgrpn"); + sym_gforce_dynamic = install("gforce_dynamic"); sym_anyna = install("anyna"); sym_anyinfnan = install("anyinfnan"); sym_anynotascii = install("anynotascii");