Skip to content

Commit fd3a72d

Browse files
venom1204ben-schwen
venom1204
and
ben-schwen
authored
Add limit parameter to nafill and setnafill (#7819)
* updated changes * added changes * .. * .. * updated changes * .. * .. * added test for that uncovered line * .. * clean up nocov * clean up ternary shenanigans * clean up nocov spill * clean up unnecessary change * clean up spill * clean up spill * remove spill * add limit for const * fix boundary case * update tests * pull var out of cases * add coverage --------- Co-authored-by: Benjamin Schwendinger <benjaminschwe@gmail.com>
1 parent ec699aa commit fd3a72d

7 files changed

Lines changed: 152 additions & 44 deletions

File tree

‎NEWS.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@
4848

4949
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.
5050

51+
14. `nafill()` and `setnafill()` gain a `limit` argument to restrict the maximum number of consecutive `NA` values filled, [#7677](https://github.com/Rdatatable/data.table/issues/7677). Thanks to @jaynewton for the suggestion and @venom1204 and @ben-schwen for the PR.
52+
5153
### BUG FIXES
5254

5355
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/shift.R‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,17 +26,21 @@ shift = function(x, n=1L, fill, type=c("lag", "lead", "shift", "cyclic"), give.n
2626
ans
2727
}
2828

29-
nafill = function(x, type=c("const","locf","nocb"), fill=NA, nan=NA) {
29+
nafill = function(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, limit=Inf) {
3030
type = match.arg(type)
31-
.Call(CnafillR, x, type, fill, nan_is_na(nan), FALSE, NULL)
31+
if (!is.numeric(limit) || length(limit) != 1L || is.na(limit) || limit < 0)
32+
stopf("limit must be a non-negative scalar numeric")
33+
.Call(CnafillR, x, type, fill, nan_is_na(nan), FALSE, NULL, as.double(floor(limit)))
3234
}
3335

34-
setnafill = function(x, type=c("const","locf","nocb"), fill=NA, nan=NA, cols=seq_along(x)) {
36+
setnafill = function(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, cols=seq_along(x), limit=Inf) {
3537
type = match.arg(type)
3638
if (is.logical(cols)) {
3739
if (length(cols) != length(x)) stopf("'cols' is a logical vector of length %d but there are %d columns", length(cols), length(x))
3840
if (anyNA(cols)) stopf("'cols' contains NA at position %d", which(is.na(cols))[1L])
3941
cols = which(cols)
4042
}
41-
invisible(.Call(CnafillR, x, type, fill, nan_is_na(nan), TRUE, cols))
43+
if (!is.numeric(limit) || length(limit) != 1L || is.na(limit) || limit < 0)
44+
stopf("limit must be a non-negative scalar numeric")
45+
invisible(.Call(CnafillR, x, type, fill, nan_is_na(nan), TRUE, cols, as.double(floor(limit))))
4246
}

‎inst/tests/nafill.Rraw‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,9 @@ if (test_bit64) {
249249
test(9.32, nafill(x, "nocb", 0), seti64(nafill(x, "nocb"), 5:6, as.integer64(0)))
250250
test(9.33, nafill(x, "locf", -1), seti64(nafill(x, "locf"), 1:2, as.integer64(-1)))
251251
test(9.34, nafill(x, "nocb", -1), seti64(nafill(x, "nocb"), 5:6, as.integer64(-1)))
252+
test(9.35, nafill(x, "const", 0, limit=1), as.integer64(c(0,NA,3,4,0,NA)))
253+
test(9.36, nafill(x, "locf", 0, limit=1), as.integer64(c(0,NA,3,4,4,NA)))
254+
test(9.37, nafill(x, "nocb", 0, limit=1), as.integer64(c(NA,3,3,4,NA,0)))
252255
}
253256

254257
# coerceAs verbose

‎inst/tests/tests.Rraw‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21985,3 +21985,29 @@ DT3 = data.table(a=c(1,NA), b=c("x",NA), c=c(3,NA))
2198521985
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)))
2198621986
test(2387.05, setnafill(copy(DT3), type="locf", cols=c(TRUE,NA,FALSE)), error="'cols' contains NA at position 2")
2198721987
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")
21988+
21989+
# add limit argument to nafill() #7677
21990+
test(2388.01, nafill(c(1, NA, NA, 5, NA, NA, NA, 9), type="locf", limit=2), c(1, 1, 1, 5, 5, 5, NA, 9))
21991+
test(2388.02, nafill(c(1, NA, NA, 4, NA), type="locf", limit=Inf), c(1, 1, 1, 4, 4))
21992+
test(2388.03, nafill(c(1L, NA, NA), type="locf", limit=1), c(1L, 1L, NA))
21993+
test(2388.04, nafill(c("a", NA, NA), type="locf", limit=1), c("a", "a", NA))
21994+
test(2388.05, nafill(c(NA, NA, 3, NA, NA), type="const", limit=1, fill=3), c(3, NA, 3, 3, NA))
21995+
test(2388.06, nafill(c(1, NA, NA, 5), type="nocb", limit=1), c(1, NA, 5, 5))
21996+
test(2388.07, setnafill(data.table(a=c(1, NA, NA)), type="locf", limit=1), data.table(a=c(1, 1, NA)))
21997+
test(2388.08, nafill(c(1, NA, NA), type="locf", limit=0), c(1, NA, NA))
21998+
test(2388.09, nafill(c(NA, 1), "locf", fill=0, limit=0), c(NA, 1))
21999+
test(2388.10, nafill(c(NA, 1, NA), type="nocb", limit=1.5), c(1, 1, NA))
22000+
test(2388.11, nafill(c(1, NA, NA), type="const", fill=0, limit=1), c(1, 0, NA))
22001+
lims = list(-1, NA, NaN, NULL, c(1, 2), 1+2i); for (i in seq_along(lims)) { test(2388.12 + i / 1000, nafill(1:5, limit=lims[[i]]), error="limit must be a non-negative scalar numeric") }
22002+
test(2388.13, setnafill(data.table(a=1:5), limit=-1), error="limit must be a non-negative scalar numeric")
22003+
test(2388.14, nafill(c(1.1, NA, NA, 4.4), type="nocb", limit=1), c(1.1, NA, 4.4, 4.4))
22004+
test(2388.15, nafill(c(1.1, NA, NA), type="locf", limit=1, nan=NaN), c(1.1, 1.1, NA))
22005+
test(2388.16, nafill(c(NA, NA, 3, NA, NA, NA, 9), type="const", fill=0, limit=2), c(0, 0, 3, 0, 0, NA, 9))
22006+
test(2388.17, nafill(c(1, NA, NA), type="const", fill=0, limit=0), c(1, NA, NA))
22007+
test(2388.18, nafill(c(NA_integer_, NA_integer_), type="const", fill=1L, limit=1), c(1L, NA))
22008+
test(2388.19, nafill(c(NA, 1L, NA), type="nocb", limit=1), c(1L, 1L, NA))
22009+
test(2388.20, nafill(c(1, NA, NA), type="const", fill=0, limit=1, nan=NaN), c(1, 0, NA))
22010+
test(2388.21, nafill(c(1, NA, NA, 5), type="nocb", limit=1, nan=NaN), c(1, NA, 5, 5))
22011+
test(2388.22, nafill(c(1L, NA, NA, 5L), type="nocb", limit=1), c(1L, NA, 5L, 5L))
22012+
test(2388.23, nafill(c("a", NA, NA), type="const", fill="x", limit=1), c("a", "x", NA))
22013+
test(2388.24, nafill(c("a", NA, NA, "z"), type="nocb", limit=1), c("a", NA, "z", "z"))

‎man/nafill.Rd‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,16 @@
1010
Fast fill missing values using constant value, \emph{last observation carried forward} or \emph{next observation carried backward}.
1111
}
1212
\usage{
13-
nafill(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA)
14-
setnafill(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, cols=seq_along(x))
13+
nafill(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, limit=Inf)
14+
setnafill(x, type=c("const", "locf", "nocb"), fill=NA, nan=NA, cols=seq_along(x), limit=Inf)
1515
}
1616
\arguments{
1717
\item{x}{ Vector, list, data.frame or data.table of logical, numeric or character columns. }
1818
\item{type}{ Character, one of \emph{"const"}, \emph{"locf"} or \emph{"nocb"}. Defaults to \code{"const"}. }
1919
\item{fill}{ Value to be used to replace missing observations. See examples. }
2020
\item{nan}{ Either \code{NaN} or \code{NA}; if the former, \code{NaN} is treated as distinct from \code{NA}, otherwise, they are treated the same during replacement. See Examples. }
2121
\item{cols}{ Numeric, character or logical vector specifying columns to be updated. A logical vector must be the same length as the number of columns in \code{x}. }
22+
\item{limit}{The maximum number of consecutive \code{NA} values to fill. Must be a non-negative scalar numeric. Default is \code{Inf}. Fractional values are truncated via \code{floor}.}
2223
}
2324
\details{
2425
Supported types are \emph{logical}, \emph{integer}, \emph{double}, \emph{character}, and \emph{factor}, as well as classes built on top of these such as \code{Date}, \code{IDate}, and \code{POSIXct}.
@@ -51,6 +52,14 @@ nafill(dt, "nocb")
5152

5253
setnafill(dt, "locf", cols=c("v2","v3"))
5354
dt
55+
56+
# limit= restricts the number of consecutive fills
57+
y = c(1, NA, NA, NA, 5)
58+
nafill(c(NA, 1), "locf", fill=0, limit=0) # Result: NA, 1 (limit=0 fills nothing)
59+
nafill(y, "locf", limit=1) # Only fills the first NA
60+
nafill(y, "locf", limit=Inf) # Fills all NAs (default)
61+
nafill(y, "const", fill=0, limit=1)
62+
5463
}
5564
\seealso{
5665
\code{\link{shift}}, \code{\link{data.table}}, \code{\link{fcoalesce}}

‎src/data.table.h‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -311,9 +311,11 @@ SEXP copyAsGrowable(SEXP x);
311311
SEXP resizeVector(SEXP x, SEXP size);
312312

313313
// nafill.c
314-
void nafillDouble(double *x, uint_fast64_t nx, unsigned int type, double fill, bool nan_is_na, ans_t *ans, bool verbose);
315-
void nafillInteger(int32_t *x, uint_fast64_t nx, unsigned int type, int32_t fill, ans_t *ans, bool verbose);
316-
SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, SEXP cols);
314+
void nafillDouble(double *x, uint_fast64_t nx, unsigned int type, double fill, bool nan_is_na, ans_t *ans, bool verbose, uint_fast64_t limit);
315+
void nafillInteger(int32_t *x, uint_fast64_t nx, unsigned int type, int32_t fill, ans_t *ans, bool verbose, uint_fast64_t limit);
316+
void nafillInteger64(int64_t *x, uint_fast64_t nx, unsigned int type, int64_t fill, ans_t *ans, bool verbose, uint_fast64_t limit);
317+
void nafillString(const SEXP *x, uint_fast64_t nx, unsigned int type, SEXP fill, ans_t *ans, bool verbose, uint_fast64_t limit);
318+
SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, SEXP cols, SEXP limit);
317319

318320
// between.c
319321
SEXP between(SEXP x, SEXP lower, SEXP upper, SEXP incbounds, SEXP NAbounds, SEXP check);

‎src/nafill.c‎

Lines changed: 97 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,170 @@
11
#include "data.table.h"
22

3-
void nafillDouble(double *x, uint_fast64_t nx, unsigned int type, double fill, bool nan_is_na, ans_t *ans, bool verbose) {
3+
void nafillDouble(double *x, uint_fast64_t nx, unsigned int type, double fill, bool nan_is_na, ans_t *ans, bool verbose, uint_fast64_t limit) {
44
double tic=0.0;
55
if (verbose)
66
tic = omp_get_wtime();
7+
uint_fast64_t fills = 0;
78
if (type==0) { // const
89
if (nan_is_na) {
910
for (uint_fast64_t i=0; i<nx; i++) {
10-
ans->dbl_v[i] = ISNAN(x[i]) ? fill : x[i];
11+
if (ISNAN(x[i])) {
12+
if (fills < limit) { ans->dbl_v[i] = fill; fills++; }
13+
else ans->dbl_v[i] = x[i];
14+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
1115
}
1216
} else {
1317
for (uint_fast64_t i=0; i<nx; i++) {
14-
ans->dbl_v[i] = ISNA(x[i]) ? fill : x[i];
18+
if (ISNA(x[i])) {
19+
if (fills < limit) { ans->dbl_v[i] = fill; fills++; }
20+
else ans->dbl_v[i] = x[i];
21+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
1522
}
1623
}
1724
} else if (type==1) { // locf
1825
if (nan_is_na) {
19-
ans->dbl_v[0] = ISNAN(x[0]) ? fill : x[0];
26+
if (ISNAN(x[0]) && limit>0) { ans->dbl_v[0] = fill; fills = 1; }
27+
else ans->dbl_v[0] = x[0];
2028
for (uint_fast64_t i=1; i<nx; i++) {
21-
ans->dbl_v[i] = ISNAN(x[i]) ? ans->dbl_v[i-1] : x[i];
29+
if (ISNAN(x[i])) {
30+
if (fills < limit) { ans->dbl_v[i] = ans->dbl_v[i-1]; fills++; }
31+
else ans->dbl_v[i] = x[i];
32+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
2233
}
2334
} else {
24-
ans->dbl_v[0] = ISNA(x[0]) ? fill : x[0];
35+
if (ISNA(x[0]) && limit>0) { ans->dbl_v[0] = fill; fills = 1; }
36+
else ans->dbl_v[0] = x[0];
2537
for (uint_fast64_t i=1; i<nx; i++) {
26-
ans->dbl_v[i] = ISNA(x[i]) ? ans->dbl_v[i-1] : x[i];
38+
if (ISNA(x[i])) {
39+
if (fills < limit) { ans->dbl_v[i] = ans->dbl_v[i-1]; fills++; }
40+
else ans->dbl_v[i] = x[i];
41+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
2742
}
2843
}
2944
} else if (type==2) { // nocb
3045
if (nan_is_na) {
31-
ans->dbl_v[nx-1] = ISNAN(x[nx-1]) ? fill : x[nx-1];
46+
if (ISNAN(x[nx-1]) && limit>0) { ans->dbl_v[nx-1] = fill; fills = 1; }
47+
else ans->dbl_v[nx-1] = x[nx-1];
3248
for (int_fast64_t i=nx-2; i>=0; i--) {
33-
ans->dbl_v[i] = ISNAN(x[i]) ? ans->dbl_v[i+1] : x[i];
49+
if (ISNAN(x[i])) {
50+
if (fills < limit) { ans->dbl_v[i] = ans->dbl_v[i+1]; fills++; }
51+
else ans->dbl_v[i] = x[i];
52+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
3453
}
3554
} else {
36-
ans->dbl_v[nx-1] = ISNA(x[nx-1]) ? fill : x[nx-1];
55+
if (ISNA(x[nx-1]) && limit>0) { ans->dbl_v[nx-1] = fill; fills = 1; }
56+
else ans->dbl_v[nx-1] = x[nx-1];
3757
for (int_fast64_t i=nx-2; i>=0; i--) {
38-
ans->dbl_v[i] = ISNA(x[i]) ? ans->dbl_v[i+1] : x[i];
58+
if (ISNA(x[i])) {
59+
if (fills < limit) { ans->dbl_v[i] = ans->dbl_v[i+1]; fills++; }
60+
else ans->dbl_v[i] = x[i];
61+
} else { ans->dbl_v[i] = x[i]; fills = 0; }
3962
}
4063
}
4164
}
4265
if (verbose)
4366
snprintf(ans->message[0], 500, _("%s: took %.3fs\n"), __func__, omp_get_wtime()-tic);
4467
}
45-
void nafillInteger(int32_t *x, uint_fast64_t nx, unsigned int type, int32_t fill, ans_t *ans, bool verbose) {
68+
void nafillInteger(int32_t *x, uint_fast64_t nx, unsigned int type, int32_t fill, ans_t *ans, bool verbose, uint_fast64_t limit) {
4669
double tic=0.0;
4770
if (verbose)
4871
tic = omp_get_wtime();
72+
uint_fast64_t fills = 0;
4973
if (type==0) { // const
5074
for (uint_fast64_t i=0; i<nx; i++) {
51-
ans->int_v[i] = x[i]==NA_INTEGER ? fill : x[i];
75+
if (x[i]==NA_INTEGER) {
76+
if (fills < limit) { ans->int_v[i] = fill; fills++; }
77+
else ans->int_v[i] = x[i];
78+
} else { ans->int_v[i] = x[i]; fills = 0; }
5279
}
5380
} else if (type==1) { // locf
54-
ans->int_v[0] = x[0]==NA_INTEGER ? fill : x[0];
81+
if (x[0]==NA_INTEGER && limit>0) { ans->int_v[0] = fill; fills = 1; }
82+
else ans->int_v[0] = x[0];
5583
for (uint_fast64_t i=1; i<nx; i++) {
56-
ans->int_v[i] = x[i]==NA_INTEGER ? ans->int_v[i-1] : x[i];
84+
if (x[i]==NA_INTEGER) {
85+
if (fills < limit) { ans->int_v[i] = ans->int_v[i-1]; fills++; }
86+
else ans->int_v[i] = x[i];
87+
} else { ans->int_v[i] = x[i]; fills = 0; }
5788
}
5889
} else if (type==2) { // nocb
59-
ans->int_v[nx-1] = x[nx-1]==NA_INTEGER ? fill : x[nx-1];
90+
if (x[nx-1]==NA_INTEGER && limit>0) { ans->int_v[nx-1] = fill; fills = 1; }
91+
else ans->int_v[nx-1] = x[nx-1];
6092
for (int_fast64_t i=nx-2; i>=0; i--) {
61-
ans->int_v[i] = x[i]==NA_INTEGER ? ans->int_v[i+1] : x[i];
93+
if (x[i]==NA_INTEGER) {
94+
if (fills < limit) { ans->int_v[i] = ans->int_v[i+1]; fills++; }
95+
else ans->int_v[i] = x[i];
96+
} else { ans->int_v[i] = x[i]; fills = 0; }
6297
}
6398
}
6499
if (verbose)
65100
snprintf(ans->message[0], 500, _("%s: took %.3fs\n"), __func__, omp_get_wtime()-tic);
66101
}
67-
void nafillInteger64(int64_t *x, uint_fast64_t nx, unsigned int type, int64_t fill, ans_t *ans, bool verbose) {
102+
void nafillInteger64(int64_t *x, uint_fast64_t nx, unsigned int type, int64_t fill, ans_t *ans, bool verbose, uint_fast64_t limit) {
68103
double tic=0.0;
69104
if (verbose)
70105
tic = omp_get_wtime();
106+
uint_fast64_t fills = 0;
71107
if (type==0) { // const
72108
for (uint_fast64_t i=0; i<nx; i++) {
73-
ans->int64_v[i] = x[i]==NA_INTEGER64 ? fill : x[i];
109+
if (x[i]==NA_INTEGER64) {
110+
if (fills < limit) { ans->int64_v[i] = fill; fills++; }
111+
else ans->int64_v[i] = x[i];
112+
} else { ans->int64_v[i] = x[i]; fills = 0; }
74113
}
75114
} else if (type==1) { // locf
76-
ans->int64_v[0] = x[0]==NA_INTEGER64 ? fill : x[0];
115+
if (x[0]==NA_INTEGER64 && limit>0) { ans->int64_v[0] = fill; fills = 1; }
116+
else ans->int64_v[0] = x[0];
77117
for (uint_fast64_t i=1; i<nx; i++) {
78-
ans->int64_v[i] = x[i]==NA_INTEGER64 ? ans->int64_v[i-1] : x[i];
118+
if (x[i]==NA_INTEGER64) {
119+
if (fills < limit) { ans->int64_v[i] = ans->int64_v[i-1]; fills++; }
120+
else ans->int64_v[i] = x[i];
121+
} else { ans->int64_v[i] = x[i]; fills = 0; }
79122
}
80123
} else if (type==2) { // nocb
81-
ans->int64_v[nx-1] = x[nx-1]==NA_INTEGER64 ? fill : x[nx-1];
124+
if (x[nx-1]==NA_INTEGER64 && limit>0) { ans->int64_v[nx-1] = fill; fills = 1; }
125+
else ans->int64_v[nx-1] = x[nx-1];
82126
for (int_fast64_t i=nx-2; i>=0; i--) {
83-
ans->int64_v[i] = x[i]==NA_INTEGER64 ? ans->int64_v[i+1] : x[i];
127+
if (x[i]==NA_INTEGER64) {
128+
if (fills < limit) { ans->int64_v[i] = ans->int64_v[i+1]; fills++; }
129+
else ans->int64_v[i] = x[i];
130+
} else { ans->int64_v[i] = x[i]; fills = 0; }
84131
}
85132
}
86133
if (verbose)
87134
snprintf(ans->message[0], 500, _("%s: took %.3fs\n"), __func__, omp_get_wtime()-tic);
88135
}
89136

90-
void nafillString(const SEXP *x, uint_fast64_t nx, unsigned int type, SEXP fill, ans_t *ans, bool verbose) {
137+
void nafillString(const SEXP *x, uint_fast64_t nx, unsigned int type, SEXP fill, ans_t *ans, bool verbose, uint_fast64_t limit) {
91138
double tic=0.0;
92139
if (verbose)
93140
tic = omp_get_wtime();
94-
if (type==0) { // const 1Code has comments. Press enter to view.
141+
uint_fast64_t fills = 0;
142+
if (type==0) { // const
95143
for (uint_fast64_t i=0; i<nx; i++) {
96-
SET_STRING_ELT(ans->char_v, i, x[i]==NA_STRING ? fill : x[i]);
144+
if (x[i]==NA_STRING) {
145+
if (fills < limit) { SET_STRING_ELT(ans->char_v, i, fill); fills++; }
146+
else SET_STRING_ELT(ans->char_v, i, x[i]);
147+
} else { SET_STRING_ELT(ans->char_v, i, x[i]); fills = 0; }
97148
}
98149
} else if (type==1) { // locf
99-
SET_STRING_ELT(ans->char_v, 0, x[0]==NA_STRING ? fill : x[0]);
150+
if (x[0]==NA_STRING && limit>0) { SET_STRING_ELT(ans->char_v, 0, fill); fills = 1; }
151+
else SET_STRING_ELT(ans->char_v, 0, x[0]);
100152
const SEXP* thisans = SEXPPTR_RO(ans->char_v); // takes out STRING_ELT from loop
101153
for (uint_fast64_t i=1; i<nx; i++) {
102-
SET_STRING_ELT(ans->char_v, i, x[i]==NA_STRING ? thisans[i-1] : x[i]);
154+
if (x[i]==NA_STRING) {
155+
if (fills < limit) { SET_STRING_ELT(ans->char_v, i, thisans[i-1]); fills++; }
156+
else SET_STRING_ELT(ans->char_v, i, x[i]);
157+
} else { SET_STRING_ELT(ans->char_v, i, x[i]); fills = 0; }
103158
}
104159
} else if (type==2) { // nocb
105-
SET_STRING_ELT(ans->char_v, nx-1, x[nx-1]==NA_STRING ? fill : x[nx-1]);
160+
if (x[nx-1]==NA_STRING && limit>0) { SET_STRING_ELT(ans->char_v, nx-1, fill); fills = 1; }
161+
else SET_STRING_ELT(ans->char_v, nx-1, x[nx-1]);
106162
const SEXP* thisans = SEXPPTR_RO(ans->char_v); // takes out STRING_ELT from loop
107163
for (int_fast64_t i=nx-2; i>=0; i--) {
108-
SET_STRING_ELT(ans->char_v, i, x[i]==NA_STRING ? thisans[i+1] : x[i]);
164+
if (x[i]==NA_STRING) {
165+
if (fills < limit) { SET_STRING_ELT(ans->char_v, i, thisans[i+1]); fills++; }
166+
else SET_STRING_ELT(ans->char_v, i, x[i]);
167+
} else { SET_STRING_ELT(ans->char_v, i, x[i]); fills = 0; }
109168
}
110169
}
111170
if (verbose)
@@ -117,7 +176,7 @@ void nafillString(const SEXP *x, uint_fast64_t nx, unsigned int type, SEXP fill,
117176
over columns of the input data. This includes handling different data types
118177
and applying the designated filling method to each column in parallel.
119178
*/
120-
SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, SEXP cols) {
179+
SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, SEXP cols, SEXP limit) {
121180
int protecti=0;
122181
const bool verbose = GetVerbose();
123182

@@ -128,6 +187,9 @@ SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, S
128187
if (verbose)
129188
tic = omp_get_wtime();
130189

190+
const double limit_d = REAL(limit)[0];
191+
const uint_fast64_t limit_n = !R_FINITE(limit_d) || limit_d >= (double)UINT_FAST64_MAX ? UINT_FAST64_MAX : (uint_fast64_t)limit_d;
192+
131193
bool copy = !LOGICAL(inplace)[0];
132194
if (!IS_TRUE_OR_FALSE(nan_is_na_arg))
133195
error(_("'%s' must be TRUE or FALSE"), "nan_is_na"); // # nocov
@@ -253,16 +315,16 @@ SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, S
253315
switch (TYPEOF(VECTOR_ELT(x, i))) {
254316
case REALSXP : {
255317
if (isInt64[i]) {
256-
nafillInteger64(i64x[i], inx[i], itype, hasFill ? ((int64_t *)fillp[i])[0] : NA_INTEGER64, &vans[i], verbose);
318+
nafillInteger64(i64x[i], inx[i], itype, hasFill ? ((int64_t *)fillp[i])[0] : NA_INTEGER64, &vans[i], verbose, limit_n);
257319
} else {
258-
nafillDouble(dx[i], inx[i], itype, hasFill ? ((double *)fillp[i])[0] : NA_REAL, nan_is_na, &vans[i], verbose);
320+
nafillDouble(dx[i], inx[i], itype, hasFill ? ((double *)fillp[i])[0] : NA_REAL, nan_is_na, &vans[i], verbose, limit_n);
259321
}
260322
} break;
261323
case LGLSXP: case INTSXP : {
262-
nafillInteger(ix[i], inx[i], itype, hasFill ? ((int32_t *)fillp[i])[0] : NA_INTEGER, &vans[i], verbose);
324+
nafillInteger(ix[i], inx[i], itype, hasFill ? ((int32_t *)fillp[i])[0] : NA_INTEGER, &vans[i], verbose, limit_n);
263325
} break;
264326
case STRSXP : {
265-
nafillString(sx[i], inx[i], itype, hasFill ? ((SEXP *)fillp[i])[0] : NA_STRING, &vans[i], verbose);
327+
nafillString(sx[i], inx[i], itype, hasFill ? ((SEXP *)fillp[i])[0] : NA_STRING, &vans[i], verbose, limit_n);
266328
} break;
267329
}
268330
}

0 commit comments

Comments
 (0)