Open up factor support

Author: MichaelChirico

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.
 
 ### Notes
 

inst/tests/nafill.Rraw

@@ -333,6 +333,23 @@ 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"))
+test(13.03, nafill(x, fill="a"), x_fill_a <- replace(x, c(1:2, 5:6, 9:10), "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"))), x_fill_a)
+test(13.10, nafill(x, fill=NA), x)
+test(13.11, nafill(x, fill=NA_integer_), x)
+test(13.12, 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);
   }