Closes #1718. Adding missing factor levels handled properly on join+u…

…pdate.

NEWS.md

@@ -205,6 +205,8 @@
 
   60. Fixed test in `onAttach()` for when `Packaged` field is missing from `DESCRIPTION`, [#1706](https://github.com/Rdatatable/data.table/issues/1706); thanks @restonslacker for BR&PR.
 
+  61. Adding missing factor levels are handled correctly in case of NAs. This affected a case of join+update operation as shown in [#1718](https://github.com/Rdatatable/data.table/issues/1718). Thanks to @daniellemccool.
+
 #### NOTES
 
   1. Updated error message on invalid joins to reflect the new `on=` syntax, [#1368](https://github.com/Rdatatable/data.table/issues/1368). Thanks @MichaelChirico.

inst/tests/tests.Rraw

@@ -8867,6 +8867,18 @@ test(1672.6, capture.output(DT[1:5, .(.I[1L], V2[1L]), by = V1]),
 TT <- as.IDate("2016-04-25")
 test(1673, class(TT + 4L), c("IDate", "Date"))
 
+# fix for #1718
+A = data.table(foo = c(1, 2, 3), bar = c(4, 5, 6))
+B = data.table(foo = c(1, 2, 3, 4, 5, 6), bar = c(NA, NA, NA, 4, 5, 6))
+A[, bar := factor(bar, levels = c(4, 5), labels = c("Boop", "Beep"), exclude = 6)]
+B[, bar := factor(bar, levels = c(4, 5, NA), labels = c("Boop", "Beep", NA), exclude = NULL)]
+test(1674.1, as.integer(B[A, bar := i.bar, on="foo"]$bar), c(1:3,1:2,NA))
+A = data.table(foo = c(1, 2, 3), bar = c(4, 5, 6))
+B = data.table(foo = c(1, 2, 3, 4, 5, 6), bar = c(NA, NA, NA, 4, 5, 6))
+A[, bar := factor(bar, levels = c(4, 5), labels = c("Boop", "Beep"), exclude = 6)]
+B[, bar := factor(bar, levels = c(4, 5), labels = c("Boop", "Beep"), exclude = 6)]
+test(1674.2, as.integer(B[A, bar := i.bar, on="foo"]$bar), c(1:2,NA,1:2,NA))
+
 ##########################
 
 # TODO: Tests involving GForce functions needs to be run with optimisation level 1 and 2, so that both functions are tested all the time.

src/assign.c

@@ -542,8 +542,9 @@ SEXP assign(SEXP dt, SEXP rows, SEXP cols, SEXP newcolnames, SEXP values, SEXP v
                                 PROTECT(addlevels = growVector(addlevels, length(addlevels)+1000));
                                 protecti++;
                             }
-                            SET_STRING_ELT(addlevels,addi,thisv);
-                            SET_TRUELENGTH(thisv,++addi+length(targetlevels));  
+                            SET_STRING_ELT(addlevels,addi++,thisv);
+                            // if-else for #1718 fix
+                            SET_TRUELENGTH(thisv, (thisv != NA_STRING) ? (addi+length(targetlevels)) : NA_INTEGER);
                         }
                         INTEGER(RHS)[j] = TRUELENGTH(thisv);
                     }