Open up factor support in nafill

NEWS.md

@@ -14,7 +14,7 @@
 
 ### NEW FEATURES
 
-1. `nafill()`, `setnafill()` extended to work on logical vectors (part of [#3992](https://github.com/Rdatatable/data.table/issues/3992)). Thanks @jangorecki for the request and @MichaelChirico for the PR.
+1. `nafill()`, `setnafill()` extended to work on logical and factor vectors (part of [#3992](https://github.com/Rdatatable/data.table/issues/3992)). Thanks @jangorecki for the request and @MichaelChirico for the PR.
 
 2. `[,showProgress=]` and `options(datatable.showProgress)` now accept an integer to control the progress bar update interval in seconds, allowing finer control over progress reporting frequency; `TRUE` uses the default 3-second interval, [#6514](https://github.com/Rdatatable/data.table/issues/6514). Thanks @ethanbsmith for the report and @ben-schwen for the PR.
 

inst/tests/nafill.Rraw

@@ -333,6 +333,30 @@ test(12.08, nafill(x, fill=NA_integer_), x)
 test(12.09, nafill(x, fill=NA_real_), x)
 test(12.10, nafill(x, fill=NaN), x)
 
+## factor input
+x = rep(NA_character_, 10L)
+x[c(3:4, 7:8)] = c("a", "b", "a", "c")
+x = as.factor(x)
+test(13.01, nafill(x, "locf"), replace(replace(x, 5:6, "b"), 9:10, "c"))
+test(13.02, nafill(x, "nocb"), replace(x, c(1:2, 5:6), "a"))
+x_fill_a = replace(x, c(1:2, 5:6, 9:10), "a")
+test(13.03, nafill(x, fill="a"), x_fill_a)
+test(13.04, nafill(x, fill=1L), x_fill_a)
+test(13.05, nafill(x, fill=1.0), x_fill_a)
+test(13.06, nafill(x, fill=factor("a")), x_fill_a)
+test(13.07, nafill(x, fill=factor("a", levels=levels(x))), x_fill_a)
+test(13.08, nafill(x, fill=factor("a", levels=c("a", "b"))), x_fill_a)
+test(13.09, nafill(x, fill=factor("a", levels=c("a", "d"))), factor(x_fill_a, levels=c("a", "b", "c", "d")))
+x_fill_d = replace(factor(x, levels = c(levels(x), "d")), c(1:2, 5:6, 9:10), "d")
+test(13.10, nafill(x, fill="d"), x_fill_d)
+test(13.11, nafill(x, fill=factor("d", levels=c("a", "b", "c", "d"))), x_fill_d)
+test(13.12, nafill(x, fill=factor("d", levels=c("d", "a", "b", "c"))), x_fill_d)
+test(13.13, nafill(x, fill=factor("d", levels=c("d", "c", "b", "a"))), x_fill_d)
+test(13.14, nafill(x, fill=factor("d", levels=c("b", "c", "d"))), x_fill_d)
+test(13.15, nafill(x, fill=NA), x)
+test(13.16, nafill(x, fill=NA_integer_), x)
+test(13.17, nafill(x, fill=NA_character_), x)
+
 ## logical
 ## character
 ## factor

man/nafill.Rd

@@ -37,6 +37,11 @@ x = c(1, NA, NaN, 3, NaN, NA, 4)
 nafill(x, "locf")
 nafill(x, "locf", nan=NaN)
 
+# works for factors
+x = gl(3, 2, 10)
+is.na(x) = 1:2
+nafill(x, "nocb")
+
 # fill= applies to any leftover NA
 nafill(c(NA, x), "locf")
 nafill(c(NA, x), "locf", fill=0)

src/nafill.c

@@ -113,7 +113,7 @@ SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, S
   if (obj_scalar) {
     if (binplace)
       error(_("'x' argument is atomic vector, in-place update is supported only for list/data.table"));
-    else if (!isReal(obj) && !isInteger(obj) && !isLogical(obj))
+    else if (!isReal(obj) && TYPEOF(obj) != INTSXP && !isLogical(obj))
       error(_("'x' argument must be logical/numeric type, or list/data.table of logical/numeric types"));
     SEXP obj1 = obj;
     obj = PROTECT(allocVector(VECSXP, 1)); protecti++; // wrap into list
@@ -124,7 +124,7 @@ SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, S
   int *icols = INTEGER(ricols);
   for (int i=0; i<length(ricols); i++) {
     SEXP this_col = VECTOR_ELT(obj, icols[i]-1);
-    if (!isReal(this_col) && !isInteger(this_col) && !isLogical(this_col))
+    if (!isReal(this_col) && TYPEOF(this_col) != INTSXP && !isLogical(this_col))
       error(_("'x' argument must be logical/numeric type, or list/data.table of logical/numeric types"));
     SET_VECTOR_ELT(x, i, this_col);
   }
@@ -218,8 +218,17 @@ SEXP nafillR(SEXP obj, SEXP type, SEXP fill, SEXP nan_is_na_arg, SEXP inplace, S
 
   if (!binplace) {
     for (R_len_t i=0; i<nx; i++) {
-      if (ANY_ATTRIB(VECTOR_ELT(x, i)))
-        copyMostAttrib(VECTOR_ELT(x, i), VECTOR_ELT(ans, i));
+      SEXP xi = VECTOR_ELT(x, i);
+      if (ANY_ATTRIB(xi)) {
+        copyMostAttrib(xi, VECTOR_ELT(ans, i));
+        if (itype == 0 && hasFill && isFactor(xi)) {
+          SEXP fillLev = PROTECT(getAttrib(VECTOR_ELT(fill, i), R_LevelsSymbol));
+          if (!R_compute_identical(PROTECT(getAttrib(xi, R_LevelsSymbol)), fillLev, 0)) {
+            setAttrib(VECTOR_ELT(ans, i), R_LevelsSymbol, fillLev);
+          }
+          UNPROTECT(2);
+        }
+      }
     }
     SEXP obj_names = getAttrib(obj, R_NamesSymbol); // copy names
     if (!isNull(obj_names)) {