diff --git a/R/unnest.R b/R/unnest.R index 98bbc63..daf1416 100644 --- a/R/unnest.R +++ b/R/unnest.R @@ -3,7 +3,8 @@ #' Quickly unnest data tables, particularly those nested by \code{dt_nest()}. #' #' @param dt_ the data table to unnest -#' @param col the column to unnest +#' @param col the column to unnest +#' @param keep whether to keep the nested column, default is \code{TRUE} #' #' @examples #' @@ -19,14 +20,14 @@ #' @import data.table #' #' @export -dt_unnest <- function(dt_, col) { +dt_unnest <- function(dt_, col, keep = TRUE) { UseMethod("dt_unnest", dt_) } #' @export -dt_unnest.default <- function(dt_, col) { - if (isFALSE(is.data.table(dt_))) { - dt_ <- as.data.table(dt_) +dt_unnest.default <- function(dt_, col, keep = TRUE) { + if (isFALSE(data.table::is.data.table(dt_))) { + dt_ <- data.table::as.data.table(dt_) } # col to unnest @@ -45,7 +46,13 @@ dt_unnest.default <- function(dt_, col) { others <- others[!others_class %in% c("list", "data.table", "data.frame", "tbl_df")] # Join them all together - dt_[seq_len(.N), eval(col)[[1L]], by = others][dt_, on = others] + dt_2 <- dt_[seq_len(.N), eval(col)[[1L]], by = others] + + if (isFALSE(keep)){ + dt_[, deparse(col) := NULL] + } + + dt_2[dt_, on = others] } diff --git a/man/dt_unnest.Rd b/man/dt_unnest.Rd index 23c5220..c02cc76 100644 --- a/man/dt_unnest.Rd +++ b/man/dt_unnest.Rd @@ -4,12 +4,14 @@ \alias{dt_unnest} \title{Unnest: Fast Unnesting of Data Tables} \usage{ -dt_unnest(dt_, col) +dt_unnest(dt_, col, keep = TRUE) } \arguments{ \item{dt_}{the data table to unnest} \item{col}{the column to unnest} + +\item{keep}{whether to keep the nested column, default is \code{TRUE}} } \description{ Quickly unnest data tables, particularly those nested by \code{dt_nest()}. diff --git a/tests/testthat/test-dt_nest.R b/tests/testthat/test-dt_nest.R index 9e23f3d..497122b 100644 --- a/tests/testthat/test-dt_nest.R +++ b/tests/testthat/test-dt_nest.R @@ -5,7 +5,7 @@ starwars <- dplyr::starwars test_that("dt_nest() works", { res <- dt_nest(starwars, species, homeworld) - expect_is(dplyr::pull(res), "list") + expect_type(dplyr::pull(res), "list") dplyr_res <- dplyr::group_nest(starwars, species, homeworld) @@ -107,6 +107,23 @@ test_that("can unnest mixture of name and unnamed lists of same length", { ) }) + +test_that("unnest keep argument", { + test_df <- data.table::data.table(x = 1:2) + nested_df <- data.table::data.table(id = 1L:2L, list_column = list(test_df, test_df)) + expect_identical( + dt_unnest(nested_df, list_column, keep = TRUE), + data.table::data.table(id = c(1L,1L,2L,2L), x = c(1L,2L,1L,2L), list_column = list(test_df, test_df)) + ) + + expect_identical( + dt_unnest(nested_df, list_column, keep = FALSE), + data.table::data.table(id = c(1L,1L,2L,2L), x = c(1L,2L,1L,2L)) + ) +}) + + + # needs to be implemented in dt_hoist() # test_that("hoist: elements must all be of same type", { # df <- data.table::data.table(x = list(1, "a"),